DistributedSampler

class braintools.trainer.DistributedSampler(data_source, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)

Sampler for distributed training.

Splits data across multiple replicas (devices/processes).

Parameters:

• data_source (Dataset) – Dataset to sample from.

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

• rank (int | None) – Rank of current replica. Default is 0.

• shuffle (bool) – Whether to shuffle the data.

• seed (int) – Random seed for shuffling.

• drop_last (bool) – Whether to drop samples that don’t divide evenly.

Examples

>>> sampler = DistributedSampler(dataset, num_replicas=4, rank=0)
>>> # Each replica gets 1/4 of the data

set_epoch(epoch)

Set the epoch for deterministic shuffling.