Data Datamodule Kfold

datamodule_kfold

CLASS	DESCRIPTION
DataModuleWeightedKFold	PyTorch Lightning DataModule with k-fold cross-validation support.

Classes

DataModuleWeightedKFold(dataset: DatasetWeightedType | dict, n_splits: int = 5, batch_size: int = 1000, shuffle: bool = True, stratify: bool = False, stratify_batches: bool = False, random_seed: int | None = None)

Bases: LightningDataModule

PyTorch Lightning DataModule with k-fold cross-validation support.

Provides stratified or random k-fold splitting for cross-validation workflows. Integrates with PyTorch Lightning Trainer for automated CV training.

PARAMETER	DESCRIPTION
dataset	Dataset instance or dictionary with "X" key (required) and optional "weights", "target" keys. TYPE: DatasetWeighted | dict
n_splits	Number of folds for cross-validation. TYPE: int, optional, by default 5 DEFAULT: 5
batch_size	Batch size for all dataloaders. TYPE: int, optional, by default 1000 DEFAULT: 1000
shuffle	Whether to shuffle data before splitting into folds. TYPE: bool, optional, by default True DEFAULT: True
stratify	Whether to use stratified splitting for folds (requires targets). TYPE: bool, optional, by default False DEFAULT: False
stratify_batches	Whether to use stratified sampling within each batch (requires targets). This is independent of fold stratification. TYPE: bool, optional, by default False DEFAULT: False
random_seed	Random seed for reproducible splits. TYPE: int | None, optional, by default None DEFAULT: None

Properties

dataset : DatasetWeighted The underlying dataset.

n_splits : int Number of folds.

current_fold : int Current fold index (0 to n_splits-1).

batch_size : int Batch size for dataloaders.

shuffle : bool Whether data is shuffled before splitting.

stratify : bool Whether stratified splitting is used for folds.

stratify_batches : bool Whether stratified sampling is used within batches.

random_seed : int | None Random seed for reproducibility.

Examples:

Basic k-fold cross-validation:

>>> import torch
>>> from spectre.data import DatasetWeighted, DataModuleWeightedKFold
>>> import pytorch_lightning as pl
>>> X = torch.randn(1000, 20)
>>> dataset = DatasetWeighted(X)
>>> kfold_dm = DataModuleWeightedKFold(dataset, n_splits=5, batch_size=128)
>>> # Train on each fold
>>> for fold in range(5):
...     kfold_dm.set_fold(fold)
...     kfold_dm.setup()
...     trainer = pl.Trainer(max_epochs=10)
...     # trainer.fit(model, kfold_dm)

Stratified k-fold for classification:

>>> X = torch.randn(1000, 20)
>>> targets = torch.randint(0, 3, (1000,))
>>> dataset = DatasetWeighted(X, target=targets)
>>> kfold_dm = DataModuleWeightedKFold(
...     dataset,
...     n_splits=5,
...     stratified=True,
...     random_seed=42,
... )
>>> for fold in range(5):
...     kfold_dm.set_fold(fold)
...     kfold_dm.setup()
...     # Training on fold with balanced class distribution

Notes

• Stratified splitting requires dataset to have targets
• Call set_fold() before setup() for each fold
• Supports all PyTorch Lightning Trainer features

METHOD	DESCRIPTION
set_fold	Set the current fold for training and validation.
get_fold_indices	Get train and validation indices for a specific fold.
setup	Setup train and validation datasets for current fold.
train_dataloader	Return dataloader for training dataset.
val_dataloader	Return dataloader for validation dataset.
check_setup	Check if the datamodule has been set up and raise error if not.

Source code in spectre/data/datamodule_kfold.py

def __init__(
    self,
    dataset: DatasetWeightedType | dict,
    n_splits: int = 5,
    batch_size: int = 1000,
    shuffle: bool = True,
    stratify: bool = False,
    stratify_batches: bool = False,
    random_seed: int | None = None,
) -> None:
    super().__init__()

    if not isinstance(n_splits, int):
        raise TypeError("`n_splits` must be an integer.")
    if n_splits < 2:
        raise ValueError("`n_splits` must be at least 2.")
    self.n_splits = n_splits

    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(shuffle, bool):
        raise TypeError("`shuffle` must be a boolean.")
    self.shuffle = shuffle

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

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

    if random_seed is not None and not isinstance(random_seed, int):
        raise TypeError("`random_seed` must be an integer or None.")
    self.random_seed = random_seed

    if isinstance(dataset, (DatasetWeighted, DatasetWeightedMemoryMapped)):
        self.dataset = dataset
    elif isinstance(dataset, dict):
        if "X" not in dataset:
            raise ValueError(
                "Key `X` must be present when passing a dictionary as dataset."
            )
        self.dataset = DatasetWeighted(
            X=dataset["X"],
            weights=dataset["weights"] if "weights" in dataset else None,
            target=dataset["target"] if "target" in dataset else None,
        )
    else:
        raise ValueError(
            f"Unsupported dataset type: {type(dataset).__name__}. Expected "
            f"`DatasetWeighted`, `DatasetWeightedMemoryMapped`, or `dict`."
        )

    # Validate stratified requirements
    if self.stratify and not self.dataset.has_target():
        raise ValueError("Stratified splitting requires dataset to have targets.")

    # Validate batch stratification requirements
    if self.stratify_batches and not self.dataset.has_target():
        raise ValueError("Batch stratification requires dataset to have targets.")

    # Generate fold indices
    self._fold_indices = self._generate_fold_indices()

    # Current fold
    self.current_fold = 0

    # Datasets and dataloaders
    self._train_dataset = None
    self._val_dataset = None
    self._dataloaders = {}
    self._setup = False

Functions

set_fold(fold_idx: int) -> None

Set the current fold for training and validation.

PARAMETER	DESCRIPTION
fold_idx	Fold index (0 to n_splits-1). TYPE: int

Source code in spectre/data/datamodule_kfold.py

def set_fold(self, fold_idx: int) -> None:
    """
    Set the current fold for training and validation.

    Parameters
    ----------
    fold_idx : int
        Fold index (0 to n_splits-1).
    """
    if not isinstance(fold_idx, int):
        raise TypeError("`fold_idx` must be an integer.")
    if fold_idx < 0 or fold_idx >= self.n_splits:
        raise ValueError(f"`fold_idx` must be in range [0, {self.n_splits - 1}].")

    self.current_fold = fold_idx
    self._setup = False  # Reset setup flag
    self._dataloaders = {}  # Clear cached dataloaders

get_fold_indices(fold_idx: int) -> tuple[torch.Tensor, torch.Tensor]

Get train and validation indices for a specific fold.

PARAMETER	DESCRIPTION
fold_idx	Fold index (0 to n_splits-1). TYPE: int

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	Train and validation indices.

Source code in spectre/data/datamodule_kfold.py

def get_fold_indices(self, fold_idx: int) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Get train and validation indices for a specific fold.

    Parameters
    ----------
    fold_idx : int
        Fold index (0 to n_splits-1).

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Train and validation indices.
    """
    if not isinstance(fold_idx, int):
        raise TypeError("`fold_idx` must be an integer.")
    if fold_idx < 0 or fold_idx >= self.n_splits:
        raise ValueError(f"`fold_idx` must be in range [0, {self.n_splits - 1}].")

    return self._fold_indices[fold_idx]

setup(stage: str | None = None) -> None

Setup train and validation datasets for current fold.

PARAMETER	DESCRIPTION
stage	PyTorch Lightning stage (unused but required by interface). TYPE: str | None, optional, by default None DEFAULT: None

Source code in spectre/data/datamodule_kfold.py

def setup(self, stage: str | None = None) -> None:
    """
    Setup train and validation datasets for current fold.

    Parameters
    ----------
    stage : str | None, optional, by default None
        PyTorch Lightning stage (unused but required by interface).
    """
    train_indices, val_indices = self._fold_indices[self.current_fold]

    self._train_dataset = Subset(self.dataset, train_indices.tolist())
    self._val_dataset = Subset(self.dataset, val_indices.tolist())

    self._setup = True

train_dataloader() -> DataLoaderWeighted

Return dataloader for training dataset.

Source code in spectre/data/datamodule_kfold.py

def train_dataloader(self) -> DataLoaderWeighted:
    """Return dataloader for training dataset."""
    self.check_setup()
    if "train" not in self._dataloaders:
        if self._train_dataset is None:
            raise RuntimeError("Training dataset is not available.")
        self._dataloaders["train"] = DataLoaderWeighted(
            dataset=self._train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            stratify=self.stratify_batches,
            drop_last=True,
        )
    return self._dataloaders["train"]

val_dataloader() -> DataLoaderWeighted

Return dataloader for validation dataset.

Source code in spectre/data/datamodule_kfold.py

def val_dataloader(self) -> DataLoaderWeighted:
    """Return dataloader for validation dataset."""
    self.check_setup()
    if "val" not in self._dataloaders:
        if self._val_dataset is None:
            raise RuntimeError("Validation dataset is not available.")
        self._dataloaders["val"] = DataLoaderWeighted(
            dataset=self._val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            stratify=False,  # Don't stratify validation batches
            drop_last=False,
        )
    return self._dataloaders["val"]

check_setup() -> None

Check if the datamodule has been set up and raise error if not.

Source code in spectre/data/datamodule_kfold.py

def check_setup(self) -> None:
    """Check if the datamodule has been set up and raise error if not."""
    if not self._setup:
        raise RuntimeError(
            "DataModule has not been set up. Call `.setup()` before accessing "
            "dataloaders."
        )

Functions