nvalchemi.data.datapipes.MultiDatasetSampler

class nvalchemi.data.datapipes.MultiDatasetSampler(dataset, *, weights=None, num_samples=None, replacement=True, shuffle=True, generator=None, num_replicas=None, rank=None, distributed_manager=None, seed=0, drop_last=False)

Sample global indices from a MultiDataset at dataset-level rates.

MultiDatasetSampler yields individual global sample indices (it is a torch.utils.data.Sampler of int), choosing which child dataset each sample is drawn from according to per-dataset weights – defaulting to the child lengths, which reproduces proportional sampling from the concatenated index space. Pass it to a DataLoader as sampler=; the loader then groups the emitted indices into batches of batch_size, so batch composition is stochastic. Use MultiDatasetBatchSampler instead when each batch must contain a guaranteed number of samples from each child.

The sampler is distributed-aware: it shards the epoch across num_replicas ranks (inferred from an initialized distributed_manager when one is supplied), and set_epoch() reseeds shuffling per epoch for correct cross-rank ordering – the same contract as torch.utils.data.DistributedSampler. replacement controls whether a child’s samples may repeat within an epoch; with replacement=False the requested per-child counts may not exceed the child sizes.

Parameters:

• dataset (MultiDataset) – Dataset wrapper that defines child dataset offsets.

• weights (Sequence[float] | None, default=None) – Per-child dataset sampling rates. None uses child lengths, matching proportional sampling from the concatenated global index space.

• num_samples (int | None, default=None) – Number of global indices emitted per epoch. None emits len(dataset) samples.

• replacement (bool, default=True) – Whether local samples may repeat within an epoch.

• shuffle (bool, default=True) – Randomize dataset choices and local sample order.

• generator (torch.Generator | None, default=None) – Optional random generator for reproducible sampling.

• num_replicas (int | None, default=None) – Number of distributed ranks. None uses initialized distributed_manager.world_size or defaults to 1.

• rank (int | None, default=None) – Rank for this sampler. None uses initialized distributed_manager.rank or defaults to 0.

• distributed_manager (DistributedManager | None, default=None) – Optional distributed manager used to infer rank and world size.

• seed (int, default=0) – Base seed used for deterministic shuffling across epochs when generator is None.

• drop_last (bool, default=False) – Drop tail samples to make the epoch evenly divisible across ranks.

Examples

Oversample a small child dataset by weighting it above its natural share:

>>> from nvalchemi.data.datapipes import DataLoader
>>> from nvalchemi.data.datapipes.samplers import MultiDatasetSampler
>>> sampler = MultiDatasetSampler(multi, weights=(1.0, 3.0))
>>> loader = DataLoader(multi, batch_size=8, sampler=sampler)

See also

MultiDatasetBatchSampler

Fix the per-child composition of every batch.

MultiDataset

The concatenated dataset these global indices address.

set_epoch(epoch)

Set the epoch used for deterministic distributed shuffling.

Parameters:

epoch (int) – Epoch number added to seed when this sampler owns its generator.

Return type:

None