Fixed issue #135; rename total_batch_size to train_batch_size (#137)

1. Remove the `subset_num_batches` from the dataset hparams. Synthetic datasets should instead use the length of the real dataset as the size, or have a configurable size
2. Add `train_subset_num_batches` and `eval_subset_num_batches` to the trainer hparams
3. Add a check in the trainer that ensures that, if this field is set, then `DatasetHparams.shuffle is False`, or otherwise emit a warning that every epoch may be using a different subset of samples
4. Renamed `total_batch_size` to `train_batch_size`. Updated hparams.

Author: ravi-mosaicml

composer/core/state.py

@@ -56,6 +56,7 @@
     "eval_dataloader",
     "precision",
     "precision_context",
+    "_steps_per_epoch",
 ]
 
 
@@ -116,6 +117,7 @@ class State(Serializable):
     # but the getter will always return a Precision enum
     precision: Union[str, types.Precision]  # type: ignore
     _precision: types.Precision = field(init=False)  # but store an enum internally
+    _steps_per_epoch: Optional[int] = field(init=False, default=None)
     precision_context: Callable[[Union[str, Precision]], ContextManager] = \
         field(default_factory=default_precision_factory)
 
@@ -210,9 +212,13 @@ def batch_idx(self) -> int:
     @property
     def steps_per_epoch(self) -> int:
         """int: The number of steps (batches) per epoch."""
-        if self.train_dataloader is None:
-            raise RuntimeError("To determine the number of steps per epoch, state.train_dataloader must be set.")
-        return len(self.train_dataloader)
+        if self._steps_per_epoch is None:
+            return len(self.train_dataloader)
+        return self._steps_per_epoch
+
+    @steps_per_epoch.setter
+    def steps_per_epoch(self, val: Optional[int]):  # type: ignore
+        self._steps_per_epoch = val
 
     @property
     def precision(self) -> types.Precision:

composer/datasets/brats.py

@@ -15,7 +15,6 @@
 from composer.datasets.dataloader import DataloaderHparams
 from composer.datasets.hparams import DatasetHparams
 from composer.utils import ddp
-from composer.utils.data import get_subset_dataset
 
 PATCH_SIZE = [1, 192, 160]
 
@@ -48,9 +47,6 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
             raise ValueError("datadir must be specified if self.synthetic is False")
         x_train, y_train, x_val, y_val = get_data_split(self.datadir)
         dataset = PytTrain(x_train, y_train, oversampling) if self.is_train else PytVal(x_val, y_val)
-        if self.subset_num_batches is not None:
-            size = batch_size * self.subset_num_batches * ddp.get_world_size()
-            dataset = get_subset_dataset(size, dataset)
         collate_fn = None if self.is_train else _my_collate
         sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
 

composer/datasets/cifar10.py

@@ -2,7 +2,6 @@
 
 from dataclasses import dataclass
 
-import torch.utils.data
 import yahp as hp
 from torchvision import transforms
 from torchvision.datasets import CIFAR10
@@ -12,7 +11,6 @@
 from composer.datasets.hparams import DatasetHparams, SyntheticHparamsMixin
 from composer.datasets.synthetic import SyntheticBatchPairDataset
 from composer.utils import ddp
-from composer.utils.data import get_subset_dataset
 
 
 @dataclass
@@ -28,20 +26,15 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
         cifar10_mean, cifar10_std = [0.4914, 0.4822, 0.4465], [0.247, 0.243, 0.261]
 
         if self.use_synthetic:
-            if self.subset_num_batches is None:
-                raise ValueError("subset_num_batches is required if use_synthetic is True")
+            total_dataset_size = 50_000 if self.is_train else 10_000
             dataset = SyntheticBatchPairDataset(
-                total_dataset_size=self.subset_num_batches * batch_size,
+                total_dataset_size=total_dataset_size,
                 data_shape=[3, 32, 32],
                 num_classes=10,
                 num_unique_samples_to_create=self.synthetic_num_unique_samples,
                 device=self.synthetic_device,
                 memory_format=self.synthetic_memory_format,
             )
-            if self.shuffle:
-                sampler = torch.utils.data.RandomSampler(dataset)
-            else:
-                sampler = torch.utils.data.SequentialSampler(dataset)
 
         else:
             if self.datadir is None:
@@ -66,10 +59,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
                 download=self.download,
                 transform=transformation,
             )
-            if self.subset_num_batches is not None:
-                size = batch_size * self.subset_num_batches * ddp.get_world_size()
-                dataset = get_subset_dataset(size, dataset)
-            sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
+        sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
 
         return dataloader_hparams.initialize_object(dataset,
                                                     batch_size=batch_size,

composer/datasets/hparams.py

@@ -90,14 +90,6 @@ class DatasetHparams(hp.Hparams, abc.ABC, metaclass=metaclass):
             If the number of samples is not divisible by the batch size, whether
             to drop the last batch (the default) or pad the last batch with zeros.
         shuffle (bool): Whether to shuffle the dataset. Defaults to True.
-        subset_num_batches (int, optional): If specified, limit the number of batches per dataloader iteration.
-            Specifically, ``len(dataloader) == num_total_batches``, where the ``dataloader`` is returned via
-            :meth:`initialize_object`. Each epoch should yield the same subset of samples.
-            
-            If this value is greater than the total number of samples in the dataset, then a :class:`ValueError` 
-            is raised.
-
-            If None (the default), then the entire dataset will be iterated over.
     """
 
     is_train: bool = hp.optional("Whether to load the training data (the default) or validation data.", default=True)
@@ -106,9 +98,6 @@ class DatasetHparams(hp.Hparams, abc.ABC, metaclass=metaclass):
                                   default=True)
     shuffle: bool = hp.optional("Whether to shuffle the dataset for each epoch. Defaults to True.", default=True)
 
-    subset_num_batches: Optional[int] = hp.optional(
-        "If not None, limit len(dataloader) to this many batches. If None (the default), then the dataloader will iterate over the entire dataset.",
-        default=None)
     datadir: Optional[str] = hp.optional("The path to the data directory", default=None)
 
     @abc.abstractmethod

composer/datasets/imagenet.py

@@ -17,7 +17,6 @@
 from composer.datasets.hparams import DataloaderSpec, DatasetHparams, SyntheticHparamsMixin
 from composer.datasets.synthetic import SyntheticBatchPairDataset
 from composer.utils import ddp
-from composer.utils.data import get_subset_dataset
 
 
 class TransformationFn:
@@ -80,10 +79,9 @@ class ImagenetDatasetHparams(DatasetHparams, SyntheticHparamsMixin):
     def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHparams) -> DataloaderSpec:
 
         if self.use_synthetic:
-            if self.subset_num_batches is None:
-                raise ValueError("subset_num_batches is required if use_synthetic is True")
+            total_dataset_size = 1_281_167 if self.is_train else 50_000
             dataset = SyntheticBatchPairDataset(
-                total_dataset_size=self.subset_num_batches * batch_size * ddp.get_world_size(),
+                total_dataset_size=total_dataset_size,
                 data_shape=[3, self.crop_size, self.crop_size],
                 num_classes=1000,
                 num_unique_samples_to_create=self.synthetic_num_unique_samples,
@@ -92,10 +90,6 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
             )
             collate_fn = None
             device_transform_fn = None
-            if self.shuffle:
-                sampler = torch.utils.data.RandomSampler(dataset)
-            else:
-                sampler = torch.utils.data.SequentialSampler(dataset)
         else:
 
             if self.is_train:
@@ -125,10 +119,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
             if self.datadir is None:
                 raise ValueError("datadir must be specified is self.synthetic is False")
             dataset = ImageFolder(os.path.join(self.datadir, split), transformation)
-            if self.subset_num_batches is not None:
-                size = batch_size * self.subset_num_batches * ddp.get_world_size()
-                dataset = get_subset_dataset(size, dataset)
-            sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
+        sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
 
         return DataloaderSpec(dataloader=dataloader_hparams.initialize_object(
             dataset=dataset,

composer/datasets/lm_datasets.py

@@ -12,7 +12,6 @@
 from composer.datasets.dataloader import DataloaderHparams
 from composer.datasets.hparams import DataloaderSpec, DatasetHparams
 from composer.utils import ddp
-from composer.utils.data import get_subset_dataset
 
 log = logging.getLogger(__name__)
 
@@ -98,10 +97,6 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
         log.info(f"Total number of samples: {num_samples:e}")
         log.info(f"Total number of tokens: {self.num_tokens:e}")
         dataset = lm_datasets
-        if self.subset_num_batches is not None:
-            size = batch_size * self.subset_num_batches * ddp.get_world_size()
-            dataset = get_subset_dataset(size, dataset)
-
         data_collator = transformers.default_data_collator
 
         sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)

composer/datasets/mnist.py

@@ -2,7 +2,6 @@
 
 from dataclasses import dataclass
 
-import torch.utils.data
 import yahp as hp
 from torchvision import datasets, transforms
 
@@ -11,7 +10,6 @@
 from composer.datasets.hparams import DatasetHparams, SyntheticHparamsMixin
 from composer.datasets.synthetic import SyntheticBatchPairDataset
 from composer.utils import ddp
-from composer.utils.data import get_subset_dataset
 
 
 @dataclass
@@ -25,20 +23,14 @@ class MNISTDatasetHparams(DatasetHparams, SyntheticHparamsMixin):
 
     def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHparams) -> DataLoader:
         if self.use_synthetic:
-            if self.subset_num_batches is None:
-                raise ValueError("subset_num_batches is required if use_synthetic is True")
             dataset = SyntheticBatchPairDataset(
-                total_dataset_size=self.subset_num_batches * batch_size,
+                total_dataset_size=60_000 if self.is_train else 10_000,
                 data_shape=[1, 28, 28],
                 num_classes=10,
                 num_unique_samples_to_create=self.synthetic_num_unique_samples,
                 device=self.synthetic_device,
                 memory_format=self.synthetic_memory_format,
             )
-            if self.shuffle:
-                sampler = torch.utils.data.RandomSampler(dataset)
-            else:
-                sampler = torch.utils.data.SequentialSampler(dataset)
 
         else:
             if self.datadir is None:
@@ -51,10 +43,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
                 download=self.download,
                 transform=transform,
             )
-            if self.subset_num_batches is not None:
-                size = batch_size * self.subset_num_batches * ddp.get_world_size()
-                dataset = get_subset_dataset(size, dataset)
-            sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
+        sampler = ddp.get_sampler(dataset, drop_last=self.drop_last, shuffle=self.shuffle)
         return dataloader_hparams.initialize_object(dataset=dataset,
                                                     batch_size=batch_size,
                                                     sampler=sampler,

composer/loggers/tqdm_logger.py

@@ -117,7 +117,9 @@ def init(self, state: State, logger: Logger) -> None:
 
     def _start(self, state: State):
         assert self.is_train is not None, "self.is_train should be set by the callback"
-        total_steps = len(state.train_dataloader) if self.is_train else len(state.eval_dataloader)
+        # TODO(anis) -- in #120, len(state.eval_dataloader) is inaccurate, as it does not incorporate
+        # trainer._eval_subset_num_batches. The evaluator spec should fix this.
+        total_steps = state.steps_per_epoch if self.is_train else len(state.eval_dataloader)
         self.pbars[self.is_train] = _TQDMLoggerInstance(total=total_steps, epoch=state.epoch, is_train=self.is_train)
 
     def epoch_start(self, state: State, logger: Logger) -> None:

composer/models/gpt2/scaling_laws_generator.py

@@ -173,7 +173,7 @@ def parse_args():
         },
     ],
     'max_epochs': 1,
-    'total_batch_size': 8,
+    'train_batch_size': 8,
     'eval_batch_size': 8,
     'seed': 17,
     'accelerator': {
@@ -320,7 +320,7 @@ def configure_mosaic_yaml(model, scaling_law_predictions):
     logger.info(f"Minumum possible serial optimization steps before SSR: {min_serial_steps:,}")
     logger.info(f"Minumum possible serial optimization steps after SSR: {math.ceil(args.ssr * min_serial_steps):,}")
     logger.info(f"Current serial optimization steps: {final_serial_steps:,}")
-    template_yaml['total_batch_size'] = batch_size
+    template_yaml['train_batch_size'] = batch_size
     assert math.floor(batch_size / curr_grad_accum) == (batch_size / curr_grad_accum)
     template_yaml['eval_batch_size'] = math.floor(batch_size / curr_grad_accum)
     template_yaml['grad_accum'] = curr_grad_accum

composer/models/unet/hparams.yaml

@@ -29,7 +29,7 @@ model:
 loggers:
   - tqdm: {}
 max_epochs: 200
-total_batch_size: 64
+train_batch_size: 64
 eval_batch_size: 8
 seed: 0
 validate_every_n_epochs: 1