Dataloader and Dataset Hyperparameter classes.


These classes are used with yahp for YAML-based configuration.


Hyperparameters to initialize a

class composer.datasets.dataset_hparams.DataLoaderHparams(num_workers=8, prefetch_factor=2, persistent_workers=True, pin_memory=True, timeout=0.0)[source]#

Bases: yahp.hparams.Hparams

Hyperparameters to initialize a

  • num_workers (int, optional) โ€“ Number of CPU workers to use per device to fetch data. Set to 0 to use the main training thread for dataloading. While zero workers can be useful for debugging, it should not be used for performance reasons. Default: 8.

  • prefetch_factor (int, optional) โ€“ Number of samples loaded in advance by each worker. For example, 2 means there will be a total of 2 * num_workers samples prefetched across all workers. If num_workers = 0, then the prefetch_factor must be left at the default value. Default: 2.

  • persistent_workers (bool) โ€“ Whether to reuse dataloader workers across epochs. If num_workers is 0, then this field must be False. Default: True.

  • pin_memory (bool, optional) โ€“ Whether or not to copy Tensors into CUDA pinned memory before returning them. If num_workers = 0, then the pin_memory must be False. Default: True.

  • timeout (float) โ€“ Timeout, in seconds, for collecting a batch from workers. Set to 0 for no timeout. Default: 0.

initialize_object(dataset, *, batch_size, sampler, drop_last, collate_fn=None, worker_init_fn=None)[source]#

Create a dataloader.

  • dataset (Dataset) โ€“ The dataset.

  • batch_size (int) โ€“ The per-device batch size.

  • sampler (Sampler[int] | None) โ€“ The sampler to use for the dataloader.

  • drop_last (bool) โ€“ Whether to drop the last batch if the number of samples is not evenly divisible by the batch size.

  • collate_fn (callable, optional) โ€“ Custom collate function. Default: None.

  • worker_init_fn (callable, optional) โ€“ Custom worker init function. Default: None.


DataLoader โ€“ The dataloader.

class composer.datasets.dataset_hparams.DatasetHparams(drop_last=True, shuffle=True)[source]#

Bases: yahp.hparams.Hparams, abc.ABC

Abstract base class for hyperparameters to initialize a dataset.

  • drop_last (bool) โ€“ If the number of samples is not divisible by the batch size, whether to drop the last batch or pad the last batch with zeros. Default: True.

  • shuffle (bool) โ€“ Whether to shuffle the dataset. Default: True.

abstract initialize_object(batch_size, dataloader_hparams)[source]#

Creates a DataLoader or DataSpec for this dataset.

  • batch_size (int) โ€“ The size of the batch the dataloader should yield. This batch size is device-specific and already incorporates the world size.

  • dataloader_hparams (DataLoaderHparams) โ€“ The dataset-independent hparams for the dataloader.

  • Iterable | DataSpec โ€“ An iterable that yields batches, or if the dataset yields batches that need custom

  • processing, a :class:`~core.data_spec.DataSpec`.