nvalchemi.data.datapipes.MultiDatasetBatchSampler

class nvalchemi.data.datapipes.MultiDatasetBatchSampler(dataset, *, batch_size, weights=None, samples_per_dataset=None, num_batches=None, epoch_policy='dataset_size', replacement=True, shuffle=True, generator=None, num_replicas=None, rank=None, distributed_manager=None, seed=0, drop_last=False)

Sample full global-index batches from a MultiDataset.

MultiDatasetBatchSampler yields whole batches – each a list of global indices (it is a torch.utils.data.Sampler of list[int]) – and is passed to a DataLoader as batch_sampler= (mutually exclusive with sampler, shuffle, and the loader’s batch_size). Unlike MultiDatasetSampler, it fixes the composition of every batch: each batch holds a deterministic number of samples from each child dataset, set either by samples_per_dataset (explicit integer counts, or floats read as relative rates) or by weights (rates that split batch_size across children). Use this class for ratio- or curriculum-controlled multi-source training where every optimizer step must see a guaranteed mixture ratio.

Epoch length (num_batches) can be given directly or derived from epoch_policy: "dataset_size" sizes the epoch by the combined length, "min_size" stops when the smallest contributing child is exhausted, and "max_size" runs until the largest is exhausted (oversampling smaller children when replacement=True). Like MultiDatasetSampler, it shards across num_replicas ranks and honors set_epoch(), mirroring torch.utils.data.DistributedSampler.

Parameters:

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

• batch_size (int) – Number of samples in each emitted batch.

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

• samples_per_dataset (Sequence[int | float] | None, default=None) – Per-child batch allocation. Integer entries are exact sample counts per batch. If any entry is a float, the full sequence is interpreted as relative per-dataset rates and allocated across batch_size. Mutually exclusive with weights.

• num_batches (int | None, default=None) – Number of batches per epoch. For replacement sampling, the default is ceil(len(dataset) / batch_size). Without replacement, the default is the number of complete batches supported by the smallest requested child allocation.

• epoch_policy ({"dataset_size", "min_size", "max_size"}, default="dataset_size") – Policy used to compute num_batches when it is not provided. "dataset_size" simply returns the combined dataset length divided by the batch size when replacement=True, otherwise min_size. "min_size" stops when the smallest contributing dataset would be exhausted. "max_size" runs until the largest contributing dataset would be exhausted, oversampling smaller datasets when replacement=True.

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

• shuffle (bool, default=True) – Randomize local sample order and sample order within each batch.

• 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 batches to make the epoch evenly divisible across ranks.

Examples

Guarantee three samples from the first child and one from the second in every batch of four:

>>> from nvalchemi.data.datapipes import DataLoader
>>> from nvalchemi.data.datapipes.samplers import MultiDatasetBatchSampler
>>> sampler = MultiDatasetBatchSampler(
...     multi, batch_size=4, samples_per_dataset=(3, 1)
... )
>>> loader = DataLoader(multi, batch_sampler=sampler)

See also

MultiDatasetSampler

Emit single indices for stochastic per-sample mixing.

MultiDataset

The concatenated dataset these global indices address.

classmethod balanced(dataset, *, batch_size, num_batches=None, epoch_policy='dataset_size', replacement=True, shuffle=True, generator=None, num_replicas=None, rank=None, distributed_manager=None, seed=0, drop_last=False)

Create a batch sampler with equal dataset-level sampling rates.

Parameters:

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

• batch_size (int) – Number of samples in each emitted batch.

• num_batches (int | None, default=None) – Number of batches per epoch.

• epoch_policy ({"dataset_size", "min_size", "max_size"}, default="dataset_size") – Policy used to compute num_batches when it is not provided.

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

• shuffle (bool, default=True) – Randomize local sample order and sample order within each batch.

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

• num_replicas (int | None, default=None) – Number of distributed ranks.

• rank (int | None, default=None) – Rank for this sampler.

• 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.

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

Returns:

Batch sampler with one equal relative weight per child dataset.

Return type:

Self

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