colossalai.utils.data_sampler

class colossalai.utils.data_sampler.DataParallelSampler(dataset, shuffle=False, seed=0, drop_last=False)[source]

A data sampler for distributed data parallelism.

Parameters

  • dataset (torch.utils.data.Dataset) – The Dataset for sampling.

  • shuffle (bool, optional) – Whether to shuffle data, defaults to False.

  • seed (int, optional) – The random seed used for sampling, defaults to 0.

  • drop_last (bool, optional) – Set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller, defaults to False.

set_epoch(epoch)[source]

Sets the epoch for this sampler. When shuffle=True, this ensures all replicas use a different random ordering for each epoch. Otherwise, the next iteration of this sampler will yield the same ordering.

Parameters

epoch (int) – Epoch number.

colossalai.utils.data_sampler.get_dataloader(dataset, shuffle=False, seed=1024, add_sampler=True, drop_last=False, pin_memory=False, num_workers=0, **kwargs)[source]

Set up a deterministic dataloader (also configure seed workers, samplers and whether shuffle or not)

Note

When pipeline parallel is enabled, shuffle cannot be True as it will result in mismatch between input data on the 1st stage and label on the last stage.

Parameters

  • dataset (torch.utils.data.Dataset) – The dataset to be loaded.

  • shuffle (bool, optional) – Whether to shuffle the dataset. Defaults to False.

  • seed (int, optional) – Random worker seed for sampling, defaults to 1024.

  • add_sampler – Whether to add DistributedDataParallelSampler to the dataset. Defaults to True.

  • drop_last (bool, optional) – Set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller, defaults to False.

  • pin_memory (bool, optional) – Whether to pin memory address in CPU memory. Defaults to False.

  • num_workers (int, optional) – Number of worker threads for this dataloader. Defaults to 0.

  • kwargs (dict) – optional parameters for torch.utils.data.DataLoader, more details could be found in DataLoader.

Returns

A DataLoader used for training or testing.

Return type

torch.utils.data.DataLoader