Data Dataloader Stratified

dataloader_stratified

CLASS	DESCRIPTION
StratifiedSampler	Sampler for stratified sampling across classes.
StratifiedBatchSampler	Batch sampler that yields stratified batches without materializing all indices.

Classes

StratifiedSampler(dataset: DatasetWeightedType, batch_size: int, shuffle: bool = True)

Bases: Sampler

Sampler for stratified sampling across classes.

Ensures balanced representation of classes in each batch by sampling proportionally from each class. Useful for imbalanced datasets.

PARAMETER	DESCRIPTION
dataset	Dataset with target labels for stratification. TYPE: DatasetWeighted
batch_size	Batch size for sampling. TYPE: int
shuffle	Whether to shuffle samples within each class. TYPE: bool, optional, by default True DEFAULT: True

Examples:

Basic usage:

>>> import torch
>>> from spectre.data import DatasetWeighted, StratifiedSampler
>>> X = torch.randn(100, 10)
>>> targets = torch.randint(0, 3, (100,))  # 3 classes
>>> dataset = DatasetWeighted(X, target=targets)
>>> sampler = StratifiedSampler(dataset, batch_size=15, shuffle=True)
>>> indices = list(sampler)
>>> len(indices)
100

Notes

• Requires dataset to have targets
• Targets should be class labels (integers)
• For continuous targets, discretize before using stratified sampling

Source code in spectre/data/dataloader_stratified.py

def __init__(
    self,
    dataset: DatasetWeightedType,
    batch_size: int,
    shuffle: bool = True,
) -> None:
    # Check if dataset has required methods (duck typing for SubsetView)
    if not (
        hasattr(dataset, "target")
        and hasattr(dataset, "has_target")
        and callable(dataset.has_target)
    ):
        if not isinstance(dataset, DatasetWeightedType):
            raise TypeError(
                "Dataset must be of type `DatasetWeighted` or `DatasetWeightedMemoryMapped`."
            )

    if not dataset.has_target():
        raise ValueError("Dataset must have targets for stratified sampling.")

    if not isinstance(batch_size, int):
        raise TypeError("Batch size must be an integer.")
    check_in_interval(batch_size, "(0, inf)")

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

    self.dataset = dataset
    self.batch_size = batch_size
    self.shuffle = shuffle

    # Get class indices
    # Convert to tensor if necessary (for DatasetWeightedMemoryMapped)
    if isinstance(dataset.target, torch.Tensor):
        targets = dataset.target.squeeze()
    else:
        # Handles np.ndarray and np.memmap
        targets = torch.from_numpy(np.array(dataset.target)).squeeze()

    self.classes, self.class_counts = torch.unique(targets, return_counts=True)

    # Create class-to-indices mapping
    self.class_indices = {}
    for cls in self.classes:
        self.class_indices[cls.item()] = (targets == cls).nonzero(as_tuple=True)[0]

    # Calculate samples per class per batch
    self.n_classes = len(self.classes)
    self.samples_per_class = max(1, batch_size // self.n_classes)

StratifiedBatchSampler(dataset: DatasetWeightedType, batch_size: int, drop_last: bool = True, shuffle: bool = True)

Bases: BatchSampler

Batch sampler that yields stratified batches without materializing all indices.

Generates batches on-the-fly with balanced class representation, avoiding memory overhead from pre-computing all sample indices. Essential for memory-mapped datasets with millions of samples where materializing indices would consume significant RAM.

PARAMETER	DESCRIPTION
dataset	Dataset with target labels for stratification. TYPE: DatasetWeightedType
batch_size	Batch size for sampling. TYPE: int
drop_last	Whether to drop the last incomplete batch. TYPE: bool, optional, by default True DEFAULT: True
shuffle	Whether to shuffle samples within each class and within each batch. TYPE: bool, optional, by default True DEFAULT: True

Examples:

Basic usage with DataLoader:

>>> import torch
>>> from torch.utils.data import DataLoader
>>> from spectre.data import DatasetWeighted, StratifiedBatchSampler
>>> X = torch.randn(1000, 10)
>>> targets = torch.randint(0, 3, (1000,))
>>> dataset = DatasetWeighted(X, target=targets)
>>> batch_sampler = StratifiedBatchSampler(dataset, batch_size=32)
>>> loader = DataLoader(dataset, batch_sampler=batch_sampler)
>>> for batch in loader:
...     print(batch.data.shape)
torch.Size([32, 10])

With multi-worker DataLoader:

>>> batch_sampler = StratifiedBatchSampler(
...     dataset,
...     batch_size=64,
...     drop_last=True,
...     shuffle=True,
... )
>>> loader = DataLoader(dataset, batch_sampler=batch_sampler, num_workers=4)

Notes

• Memory-efficient: Does not materialize all indices upfront
• Compatible with PyTorch's multi-worker DataLoader
• Requires dataset to have targets
• Targets should be discrete class labels
• For continuous targets, discretize before using stratified sampling

References

• [1] https://discuss.pytorch.org/t/how-to-do-a-stratified-split/62290

Source code in spectre/data/dataloader_stratified.py

def __init__(
    self,
    dataset: DatasetWeightedType,
    batch_size: int,
    drop_last: bool = True,
    shuffle: bool = True,
) -> None:
    if not (
        hasattr(dataset, "target")
        and hasattr(dataset, "has_target")
        and callable(dataset.has_target)
    ):
        if not isinstance(dataset, DatasetWeightedType):
            raise TypeError(
                "Dataset must be of type `DatasetWeighted` or "
                "`DatasetWeightedMemoryMapped`."
            )

    if not dataset.has_target():
        raise ValueError("Dataset must have targets for stratified sampling.")

    if not isinstance(batch_size, int):
        raise TypeError("Batch size must be an integer.")
    check_in_interval(batch_size, "(0, inf)")

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

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

    self.dataset = dataset
    self.batch_size = batch_size
    self.drop_last = drop_last
    self.shuffle = shuffle

    # Get class indices (only store indices, not all samples)
    # Convert to tensor if necessary (for DatasetWeightedMemoryMapped)
    if isinstance(dataset.target, torch.Tensor):
        targets = dataset.target.squeeze()
    else:
        # Handles np.ndarray and np.memmap
        targets = torch.from_numpy(np.array(dataset.target)).squeeze()

    self.classes, self.class_counts = torch.unique(targets, return_counts=True)

    # Create class-to-indices mapping
    self.class_indices = {}
    for cls in self.classes:
        self.class_indices[cls.item()] = (targets == cls).nonzero(as_tuple=True)[0]

    # Calculate samples per class per batch
    self.n_classes = len(self.classes)
    self.samples_per_class = max(1, batch_size // self.n_classes)

    # Calculate total batches
    if self.drop_last:
        # Minimum count determines how many complete batches we can make
        min_count = self.class_counts.min().item()
        max_batches_per_class = min_count // self.samples_per_class
        self._num_batches = max_batches_per_class
    else:
        # Maximum count determines total batches
        max_count = self.class_counts.max().item()
        max_batches_per_class = (
            max_count + self.samples_per_class - 1
        ) // self.samples_per_class
        self._num_batches = max_batches_per_class

Functions