afnio.utils.data.dataloader#

Functions

collate_tuple(items)

Recursively collates a batch of tuples, preserving nested structure.

Classes

DataLoader(dataset[, batch_size, shuffle, ...])

Data loader combines a dataset and a sampler, and provides an iterable over the given dataset.
class afnio.utils.data.dataloader.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, drop_last=False, seed=None)[source]#

Bases: Generic[T_co]

Data loader combines a dataset and a sampler, and provides an iterable over the given dataset.

The DataLoader supports both map-style and iterable-style datasets with single-process loading, customizing loading order and optional automatic batching (collation) and memory pinning.

See afnio.utils.data documentation page for more details.

Parameters:

  • dataset (Dataset) – dataset from which to load the data.

  • batch_size (int, optional) – how many samples per batch to load (default: 1).

  • shuffle (bool, optional) – set to True to have the data reshuffled at every epoch (default: False).

  • sampler (Sampler or Iterable, optional) – defines the strategy to draw samples from the dataset. Can be any Iterable with __len__ implemented. If specified, shuffle must not be specified.

  • drop_last (bool, optional) – set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of dataset is not divisible by the batch size, then the last batch will be smaller. (default: False)

  • seed (int, optional) – If not None, this seed will be used by RandomSampler to generate random indexes. (default: None)

batch_size: Optional[int]#
dataset: Dataset[TypeVar(T_co, covariant=True)]#
drop_last: bool#
sampler: Union[Sampler, Iterable]#
afnio.utils.data.dataloader.collate_tuple(items)[source]#

Recursively collates a batch of tuples, preserving nested structure.

This function should only be called when processing batches where each element is a tuple (i.e., when the dataset’s __getitem__ returns tuples).

The function first transposes the batch, so that each position in the tuple is grouped together. For each group:

  • If all elements are Variables, returns a single Variable whose .data is a list

    of the original .data fields, and whose role and requires_grad are taken from the first Variable.

  • If all elements are tuples, recursively collates them to preserve nested

    structure.

  • If some elements are tuples and some are not, recursively collates the tuples and

    leaves other elements as is, preserving their position.

  • Otherwise, returns a list of the grouped items.

This enables flexible batching for datasets that return tuples of Variables, nested tuples, or mixed structures.