nvalchemi.data.DataLoader

class nvalchemi.data.DataLoader(dataset, *, batch_size=1, shuffle=False, drop_last=False, sampler=None, batch_sampler=None, prefetch_factor=2, num_streams=4, use_streams=True, pin_memory=False, batch_transforms=None)

Batch-iterating data loader that yields Batch.

DataLoader is the consumer end of the pipeline (Reader -> Dataset (-> MultiDataset) -> DataLoader). It wraps a batch-loadable dataset – a Dataset, a MultiDataset, or any object implementing BatchDatasetProtocol – and yields graph-collated Batch objects built via from_data_list().

Sampling follows PyTorch’s conventions. With no sampler, shuffle selects a RandomSampler or SequentialSampler; a custom sampler (yielding sample indices) overrides shuffle; and a batch_sampler (yielding whole lists of indices) sets the batch composition itself and is mutually exclusive with sampler, shuffle, and batch_size. To mix several datasets, pair a MultiDataset with MultiDatasetSampler (as sampler=, per-sample rates) or MultiDatasetBatchSampler (as batch_sampler=, a fixed per-batch mixture); under distributed training DDPHook swaps in the rank-sharded variants automatically.

Compared with torch.utils.data.DataLoader, this loader yields Batch graphs rather than default-collated tensors, with graph-aware collation. Instead of forking worker processes it prefetches on background threads and overlaps device transfers on CUDA streams, and it fuses reads: prefetch_factor emitted batches are pulled from the backend in one windowed read (effective window batch_size * prefetch_factor), which amortizes I/O far better than one __getitem__ per sample. Set prefetch_factor=0 to read a single emitted batch at a time, and pin_memory=True to request page-locked tensors from readers that support it.

Parameters:

• dataset (BatchDatasetProtocol) – AtomicData-native dataset to load from.

• batch_size (int, default=1) – Number of samples per batch.

• shuffle (bool, default=False) – Randomize sample order each epoch.

• drop_last (bool, default=False) – Drop the last incomplete batch.

• sampler (torch.utils.data.Sampler | None, default=None) – Custom sampler (overrides shuffle).

• batch_sampler (torch.utils.data.Sampler | None, default=None) – Custom sampler that yields batches of sample indices.

• prefetch_factor (int, default=2) – Number of emitted batches to fuse into each backend read. The effective read window is batch_size * prefetch_factor. Set to 0 to disable fused prefetching and read one emitted batch at a time.

• num_streams (int, default=4) – Number of CUDA streams for prefetching.

• use_streams (bool, default=True) – Enable CUDA-stream prefetching.

• pin_memory (bool, default=False) – If True, request page-locked CPU tensors from readers that support pinned-memory reads.

• batch_transforms (Sequence[BatchTransform] | None, default=None) – Optional per-batch transforms applied to each yielded Batch after collation. None or an empty sequence disables the hook (zero runtime overhead on the hot path). See the Notes section for thread placement and CUDA-stream semantics. For per-sample transforms applied before collation, see Dataset (transforms parameter).

dataset

The underlying dataset.

Type:

BatchDatasetProtocol

batch_size

Number of samples per batch.

Type:

int

sampler

Resolved sampler (RandomSampler if shuffle=True, else SequentialSampler; user-supplied sampler overrides both).

Type:

torch.utils.data.Sampler

drop_last

Whether the trailing partial batch is dropped.

Type:

bool

prefetch_factor

Configured prefetch depth (see __iter__()).

Type:

int

num_streams

Configured CUDA-stream pool size for prefetching.

Type:

int

use_streams

Whether stream-based prefetching is actually enabled. Stored as use_streams and torch.cuda.is_available(); reflects runtime availability, not the raw argument.

Type:

bool

pin_memory

Whether page-locked CPU tensors are requested from compatible readers.

Type:

bool

Raises:

• ValueError – Raised at construction if batch_size < 1 or prefetch_factor < 0.

• TypeError – Raised at construction if batch_transforms is not a Sequence (e.g. a single callable or a generator was passed).

• RuntimeError – Raised during iteration (not construction) when any batch transform fails; the original exception is chained via __cause__.

Parameters:

• dataset (BatchDatasetProtocol)

• batch_size (int)

• shuffle (bool)

• drop_last (bool)

• sampler (Sampler | None)

• batch_sampler (Sampler[Sequence[int]] | None)

• prefetch_factor (int)

• num_streams (int)

• use_streams (bool)

• pin_memory (bool)

• batch_transforms (Sequence[BatchTransform] | None)

Notes

Batch transforms run on the consumer (main) thread after collation, not on the prefetch workers; the fully assembled Batch does not exist until the main thread constructs it. Transforms are applied in order via Compose and execute on the current CUDA stream at yield time; wrap iteration in your own torch.cuda.stream(...) context to control placement.

Examples

>>> from nvalchemi.data.datapipes import AtomicDataZarrReader, Dataset, DataLoader
>>> reader = AtomicDataZarrReader("dataset.zarr")
>>> ds = Dataset(reader, device="cpu")
>>> def center_positions(batch):
...     batch.positions = batch.positions - batch.positions.mean(0)
...     return batch
>>> loader = DataLoader(ds, batch_size=4, batch_transforms=[center_positions])
>>> for batch in loader:
...     print(batch.positions.shape)

property effective_read_window: int

Return the maximum sample count in one fused backend read.

set_epoch(epoch)

Set the epoch for the sampler (used in distributed training).

Parameters:

epoch (int) – Current epoch number.

Return type:

None

set_epoch_step(step)

Seek the next iterator to an intra-epoch batch offset.

Parameters:

step (int) – Number of complete batches to skip in sampler order before the next iterator starts yielding. The skip advances only the sampler/index stream; it does not load or collate skipped batches.

Raises:

ValueError – If step is negative.

Return type:

None