DistributedDataLoader

class braintools.trainer.DistributedDataLoader(dataset, batch_size=32, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False, **kwargs)

DataLoader for distributed training.

Automatically shards data across multiple devices/processes.

Parameters:

• dataset (Dataset | Any) – Dataset to load from.

• batch_size (int) – Batch size per replica.

• num_replicas (int | None) – Number of replicas. Default is number of JAX devices.

• rank (int | None) – Rank of this replica.

• shuffle (bool) – Whether to shuffle data.

• seed (int) – Random seed for shuffling.

• drop_last (bool) – Whether to drop incomplete batches.

• **kwargs – Additional arguments passed to DataLoader.

Examples

>>> loader = DistributedDataLoader(dataset, batch_size=32)
>>> # Data is automatically split across devices

set_epoch(epoch)

Set the epoch for the distributed sampler.