Data Dataloader

dataloader

CLASS	DESCRIPTION
DataLoaderWeighted	Dataloader for datasets with advanced indexing and multi-worker support.

Classes

DataLoaderWeighted(dataset: DatasetType, batch_size: int = 1000, shuffle: bool = True, drop_last: bool = True, stratify: bool = False, num_workers: int = 0, pin_memory: bool = False, persistent_workers: bool = False, prefetch_factor: int | None = None, collate_fn: Callable | None = None)

Dataloader for datasets with advanced indexing and multi-worker support.

Supports multiple PyTorch dataset types with optimized tensor-based indexing. Provides flexible batching with shuffle, multi-worker loading, and device pinning for efficient data loading in machine learning workflows.

PARAMETER	DESCRIPTION
dataset	Dataset that supports tensor and list indexing. TYPE: DatasetType
batch_size	Batch size for data loading. TYPE: int, optional, by default 1000 DEFAULT: 1000
shuffle	Whether to shuffle dataset indices. TYPE: bool, optional, by default True DEFAULT: True
drop_last	Whether to drop the last incomplete batch. TYPE: bool, optional, by default True DEFAULT: True
stratify	Whether to use stratified sampling for balanced class representation. Requires dataset to be DatasetWeighted or DatasetWeightedMemoryMapped with target labels. Each batch will contain proportional representation from each class. TYPE: bool, optional, by default False DEFAULT: False
num_workers	Number of subprocesses for data loading. 0 means data loaded in main process. TYPE: int, optional, by default 0 DEFAULT: 0
pin_memory	If True, tensors copied to CUDA pinned memory before returning them. Useful when using GPU for faster host-to-device transfers. TYPE: bool, optional, by default False DEFAULT: False
persistent_workers	If True, keeps workers alive between epochs. Only valid when num_workers > 0. TYPE: bool, optional, by default False DEFAULT: False
prefetch_factor	Number of batches loaded in advance by each worker. Only valid when num_workers > 0. If None, defaults to 2 when num_workers > 0. TYPE: int | None, optional, by default None DEFAULT: None
collate_fn	Custom function to merge list of samples into batch. If None, uses default collation. TYPE: Callable | None, optional, by default None DEFAULT: None

Properties

dataset : DatasetType The underlying dataset.

batch_size : int Current batch size.

dataset_len : int Length of the dataset.

Examples:

Basic usage with different dataset types:

>>> import torch
>>> from spectre.data import DataLoaderWeighted, DatasetWeighted
>>> from torch.utils.data import TensorDataset
>>> data = torch.randn(100, 10)
>>> weights = torch.rand(100)
>>> dataset_w = DatasetWeighted(data, weights)
>>> dataloader_w = DataLoaderWeighted(dataset_w, batch_size=20, shuffle=True)
>>> for batch_data, batch_weights in dataloader_w:
...     print(batch_data.shape, batch_weights.shape)
torch.Size([20, 10]) torch.Size([20])
torch.Size([20, 10]) torch.Size([20])

Example with TensorDataset:

>>> tensor_data = torch.randn(100, 10)
>>> tensor_dataset = TensorDataset(tensor_data)
>>> tensor_dataloader = DataLoaderWeighted(
...     tensor_dataset,
...     batch_size=15,
...     shuffle=False,
... )
>>> for batch in tensor_dataloader:
...     print(batch[0].shape)
torch.Size([15, 10])
torch.Size([15, 10])

Example with Subset:

>>> full_data = torch.randn(50, 5)
>>> full_dataset = TensorDataset(full_data)
>>> subset_indices = list(range(30))
>>> subset_dataset = Subset(full_dataset, subset_indices)
>>> subset_dataloader = DataLoaderWeighted(
...     subset_dataset,
...     batch_size=10,
...     drop_last=False,
... )
>>> for batch in subset_dataloader:
...     print(batch[0].shape)
torch.Size([10, 5])
torch.Size([10, 5])
torch.Size([10, 5])
torch.Size([0, 5])

ATTRIBUTE	DESCRIPTION
dataset	The underlying dataset. TYPE: DatasetType
batch_size	Current batch size. TYPE: int
dataset_len	Length of the dataset. TYPE: int

Source code in spectre/data/dataloader.py

def __init__(
    self,
    dataset: DatasetType,
    batch_size: int = 1000,
    shuffle: bool = True,
    drop_last: bool = True,
    stratify: bool = False,
    num_workers: int = 0,
    pin_memory: bool = False,
    persistent_workers: bool = False,
    prefetch_factor: int | None = None,
    collate_fn: Callable | None = None,
) -> None:
    if not isinstance(batch_size, int):
        raise TypeError("Batch size must be an integer.")
    check_in_interval(batch_size, "(0, inf)")
    self._batch_size = batch_size

    if not isinstance(dataset, DatasetType):
        raise TypeError(
            "Dataset must be of type `DatasetWeighted`, `DatasetWeightedMemoryMapped`, "
            "`TensorDataset`, `Subset`, or `Dataset`."
        )
    self._dataset = dataset

    if not isinstance(shuffle, bool):
        raise TypeError("`shuffle` must be a boolean.")
    self.shuffle = shuffle

    if not isinstance(drop_last, bool):
        raise TypeError("`drop_last` must be a boolean.")
    self.drop_last = drop_last

    if not isinstance(stratify, bool):
        raise TypeError("`stratify` must be a boolean.")
    if stratify:
        # Check the underlying dataset if it's a Subset
        check_dataset = dataset.dataset if isinstance(dataset, Subset) else dataset
        if not isinstance(
            check_dataset, (DatasetWeighted, DatasetWeightedMemoryMapped)
        ):
            raise TypeError(
                "Stratification requires dataset to be `DatasetWeighted` or "
                "`DatasetWeightedMemoryMapped`."
            )
        if not check_dataset.has_target():
            raise ValueError("Stratification requires dataset to have targets.")
    self.stratify = stratify

    if not isinstance(num_workers, int):
        raise TypeError("`num_workers` must be an integer.")
    if num_workers < 0:
        raise ValueError("`num_workers` must be non-negative.")
    self.num_workers = num_workers

    if not isinstance(pin_memory, bool):
        raise TypeError("`pin_memory` must be a boolean.")
    self.pin_memory = pin_memory

    if not isinstance(persistent_workers, bool):
        raise TypeError("`persistent_workers` must be a boolean.")
    if persistent_workers and num_workers == 0:
        raise ValueError("`persistent_workers` requires `num_workers > 0`.")
    self.persistent_workers = persistent_workers

    if prefetch_factor is not None:
        if not isinstance(prefetch_factor, int):
            raise TypeError("`prefetch_factor` must be an integer or None.")
        check_in_interval(prefetch_factor, "[1, inf)")
        if num_workers == 0:
            raise ValueError("`prefetch_factor` requires `num_workers > 0`.")
    self.prefetch_factor = prefetch_factor

    if collate_fn is not None and not callable(collate_fn):
        raise TypeError("`collate_fn` must be callable or None.")
    self.collate_fn = collate_fn

    # Use PyTorch DataLoader for multi-worker support
    self._use_torch_dataloader = num_workers > 0
    self._torch_dataloader: TorchDataLoader | None = None
    self.idx: torch.Tensor | None = None

Attributes

dataset: DatasetType property writable

The underlying dataset.

batch_size: int property writable

Current batch size.

dataset_len: int property

Length of the dataset.

SubsetView(target_data, indices)

Temporary workaround view for stratified sampling on Subset datasets.

Source code in spectre/data/dataloader.py

def __init__(self, target_data, indices):
    self.target = target_data
    self._indices = indices

Functions