DistributedSamplerWrapper

class core.samplers.DistributedSamplerWrapper(sampler, num_replicas=None, rank=None)

Bases: torch.utils.data.sampler.Sampler

Wrapper for distributing any sampler across multiple processes.

This wrapper takes an existing sampler and distributes its indices across multiple replicas (processes) in a distributed training setup. It ensures each replica gets a unique subset of the data by splitting the indices evenly across replicas.

Note: This wrapper supports samplers that may return a different number of samples each epoch, as it recomputes the indices on each iteration.

Note: If the length of the sampler is not a multiple of num_replicas, some samples will be dropped to ensure equal distribution across replicas.

Parameters:

• sampler (Sampler) – The base sampler to wrap

• num_replicas (Optional[int]) – Number of processes in distributed training

• rank (Optional[int]) – Rank of the current process

set_epoch(epoch)

Pass the epoch to the underlying sampler if it supports it.