omniloader.sampling.subsamplers

Per-dataset index draws with configurable reuse and exhaustion policies.

An IndexPool produces the local indices to visit within a single dataset for a given epoch. It captures three orthogonal choices:

• replacement — whether a draw may repeat an index within one epoch;

• exhaustion policy — for without-replacement draws, whether each epoch is a fresh random subsample ("fresh") or whether the pool shuffles through the entire dataset before any index is reused across epochs ("exhaust");

• weights — optional per-index weights for skewed sampling (e.g. class balancing within a dataset).

Draws are fully determined by a base seed and the epoch number, so runs are reproducible.

Classes

ExhaustionPolicy	How without-replacement draws relate across epochs.
IndexPool	Draw local indices for one dataset, honouring reuse and exhaustion policy.

Module Contents

class omniloader.sampling.subsamplers.ExhaustionPolicy(*args, **kwds)

Bases: enum.Enum

How without-replacement draws relate across epochs.

FRESH

Every epoch independently samples without replacement from the full dataset; nothing is remembered between epochs.

EXHAUST

Maintain a running shuffled ordering of the whole dataset and consume it across epochs, reshuffling only once every index has been used. This guarantees uniform coverage before any reuse.

class omniloader.sampling.subsamplers.IndexPool(size: int, replacement: bool = False, policy: ExhaustionPolicy = ExhaustionPolicy.FRESH, seed: int = 0, weights: collections.abc.Sequence[float] | None = None)

Draw local indices for one dataset, honouring reuse and exhaustion policy.

Parameters:

• size – Number of samples in the dataset.

• replacement – When True, indices may repeat within an epoch and the exhaustion policy is ignored. Defaults to False.

• policy – The ExhaustionPolicy for without-replacement draws.

• seed – Base seed; combined with the epoch to make draws reproducible.

• weights – Optional per-index sampling weights of length size (e.g. class-balancing weights). When set, draws are weighted via torch.multinomial and the exhaustion policy is bypassed.

Raises:

ValueError – If size is not positive or weights has a bad length.

draw(count: int, epoch: int) → list[int]

Return count local indices to visit this epoch.

Parameters:

• count – Number of indices requested. May exceed size when replacement is True or under the EXHAUST policy (which wraps across reshuffles).

• epoch – The epoch number, used to seed the draw.

Returns:

A list of count indices in [0, size).