nshtrainer 0.34.0__py3-none-any.whl → 0.34.2__py3-none-any.whl

nshtrainer/config.py

@@ -1,6 +1,3 @@
-from nshconfig._config import Config as Config
-from nshsnap._config import SnapshotConfig as SnapshotConfig
-
 from nshtrainer._checkpoint.loader import (
     BestCheckpointStrategyConfig as BestCheckpointStrategyConfig,
 )
@@ -65,13 +62,13 @@ from nshtrainer.callbacks.throughput_monitor import (
 )
 from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
 from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
-from nshtrainer.config import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
 from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
 from nshtrainer.loggers.tensorboard import (
     TensorboardLoggerConfig as TensorboardLoggerConfig,
 )
 from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
+from nshtrainer.lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,

nshtrainer/data/balanced_batch_sampler.py

@@ -81,6 +81,22 @@ class BalancedBatchSampler(BatchSampler):
     ):
         super().__init__(sampler, batch_size, drop_last=drop_last)
 
+        # Validate the dataset
+        dataset = self._unwrap_dataset(self.distributed_sampler.dataset)
+        # Dataset much either implement `data_sizes`, or we need to provide a custom
+        # implementation of the dataset sizes function.
+        if isinstance(dataset, DatasetWithSizes):
+            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
+
+        elif self._data_sizes_fn is not None:
+            log.critical("BalancedBatchSampler: Using custom data_sizes_fn")
+        else:
+            raise ValueError(
+                "Dataset must implement the `data_sizes` method, "
+                "or a custom data_sizes_fn must be provided "
+                "to the BalancedBatchSampler."
+            )
+
         self._device = device
         self._data_sizes_fn = data_sizes_fn
 
@@ -97,7 +113,6 @@ class BalancedBatchSampler(BatchSampler):
         # Dataset much either implement `data_sizes`, or we need to provide a custom
         # implementation of the dataset sizes function.
         if isinstance(dataset, DatasetWithSizes):
-            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
             return dataset.data_sizes(indices)
 
         if (data_sizes_fn := self._data_sizes_fn) is not None:

{nshtrainer-0.34.0.dist-info → nshtrainer-0.34.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.34.0
+Version: 0.34.2
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.34.0.dist-info → nshtrainer-0.34.2.dist-info}/RECORD

@@ -30,9 +30,9 @@ nshtrainer/callbacks/shared_parameters.py,sha256=fqlDweFDXPV_bfcAWpRgaJIad9i5Aeh
 nshtrainer/callbacks/throughput_monitor.py,sha256=H_ocXErZxUO3dxFk8Tx_VQdpI9E_Ztvqof5WtFevLyQ,1838
 nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50,4607
 nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsiHZk1-Y,2921
-nshtrainer/config.py,sha256=skar_Wfz50_sU2NZS8PEjqofWeon4g4cyIgby3Da81g,8308
+nshtrainer/config.py,sha256=6U7B-kCIMrfEnF_y92RuBm1WfASW7k05Zsm2uHBzRrk,8205
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
-nshtrainer/data/balanced_batch_sampler.py,sha256=WAjhbO9EsZ_UadhdW3obBsjvEDMc2V-irpjegqIb7AI,4791
+nshtrainer/data/balanced_batch_sampler.py,sha256=ybMJF-CguaZ17fLEweZ5suaGOiHOMEm3Bn8rQfGTzGQ,5445
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
 nshtrainer/ll/__init__.py,sha256=L-aTi1V1bbvnZjOro8NvI393zbHQSFR9movWSRK9Mds,2477
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
@@ -97,6 +97,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.34.0.dist-info/METADATA,sha256=GYC9ejdKV3MCyOFhJcFjI-uedTWLGWj-SE5S79ruug4,916
-nshtrainer-0.34.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.34.0.dist-info/RECORD,,
+nshtrainer-0.34.2.dist-info/METADATA,sha256=DQyYTUO0wpboH1gy3nSRJV6EsWCpY7Kb_ldD8v4BQFY,916
+nshtrainer-0.34.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.34.2.dist-info/RECORD,,