omniloader.sampling.sampler

A Sampler that turns a mixing strategy into a per-epoch index stream.

OmniSampler is the bridge between a MixingStrategy and PyTorch’s DataLoader. Each epoch it asks the strategy for the global indices to visit; call OmniSampler.set_epoch() before every epoch to advance the reproducible draw.

It is distributed-aware: when num_replicas/rank are given (or inferred from an initialized torch.distributed process group) it shards the epoch’s indices across replicas so each rank sees a disjoint slice, mirroring torch.utils.data.DistributedSampler.

Classes

OmniSampler

Yield global indices for one epoch according to a mixing strategy.

Module Contents

class omniloader.sampling.sampler.OmniSampler(strategy: omniloader.sampling.strategies.MixingStrategy, epoch: int = 0, num_replicas: int | None = None, rank: int | None = None, drop_last: bool = False)

Bases: torch.utils.data.Sampler[int]

Yield global indices for one epoch according to a mixing strategy.

Parameters:

• strategy – The mixing strategy that computes the epoch’s index list.

• epoch – Initial epoch number.

• num_replicas – Number of distributed replicas. None infers from torch.distributed (or 1 when not distributed).

• rank – This process’s replica index. None infers from torch.distributed (or 0).

• drop_last – When sharding, drop the tail so every rank gets exactly the same number of samples. When False (default), the index list is padded by wrapping so no sample is skipped.

set_epoch(epoch: int) → None

Set the epoch used to seed the next iteration.

Parameters:

epoch – The upcoming epoch number.

__iter__() → collections.abc.Iterator[int]

Iterate this rank’s global indices for the current epoch.

__len__() → int

Number of indices this rank produces per epoch.