Data Datamodule

datamodule

CLASS	DESCRIPTION
DataModuleWeighted	PyTorch Lightning DataModule for weighted datasets with flexible splitting.

Classes

DataModuleWeighted(dataset: DatasetWeightedType | dict, batch_size: int = 1000, train_stride: int = 1, val: float = 0.2, use_random_split: bool = True, shuffle: bool = False, stratify: bool = False)

Bases: LightningDataModule

PyTorch Lightning DataModule for weighted datasets with flexible splitting.

Provides comprehensive data management for machine learning workflows with support for weighted datasets, optional targets, flexible train/validation splits, and optimized dataloader caching.

PARAMETER	DESCRIPTION
dataset	Dataset instance or dictionary with "X" key (required) and optional "weights", "target" keys. Subset is not supported to avoid double indexing. TYPE: DatasetWeighted | dict
batch_size	Batch size for all dataloaders. TYPE: int, optional, by default 1000 DEFAULT: 1000
train_stride	Training stride for subsampling data (1 means no subsampling, all data used for training). TYPE: int, optional, by default 1 DEFAULT: 1
val	Validation split fraction in range [0.0, 1.0]. Set to 0.0 to disable validation. TYPE: float, optional, by default 0.2 DEFAULT: 0.2
use_random_split	Whether to use random data split (True) or sequential split (False). TYPE: bool, optional, by default True DEFAULT: True
shuffle	Whether to shuffle data in dataloaders. TYPE: bool, optional, by default False DEFAULT: False
stratify	Whether to stratify splits based on targets. Requires dataset to have targets. TYPE: bool, optional, by default False DEFAULT: False

Properties

dataset : DatasetWeighted The underlying dataset.

dataset_strided : Subset The strided (subsampled) dataset.

batch_size : int Current batch size.

val : float Validation split fraction.

use_random_split : bool Whether using random or sequential splitting.

shuffle : bool Whether dataloaders shuffle data.

train_stride : int Training stride for subsampling data.

train_dataloader : DataLoaderWeighted Dataloader for training dataset.

val_dataloader : DataLoaderWeighted | None Dataloader for validation dataset, or None if no validation split.

test_dataloader : DataLoaderWeighted Dataloader for testing on full dataset.

predict_dataloader : DataLoaderWeighted Dataloader for prediction on full dataset.

Examples:

Basic usage with DatasetWeighted:

>>> from spectre.data import DataModuleWeighted, DatasetWeighted
>>> import torch
>>> X = torch.randn(1000, 10)
>>> weights = torch.rand(1000)
>>> target = torch.randint(0, 2, (1000,))
>>> dataset = DatasetWeighted(X=X, weights=weights, target=target)
>>> datamodule = DataModuleWeighted(
...     dataset=dataset,
...     batch_size=128,
...     train_stride=2,
...     val=0.2,
...     use_random_split=True,
...     shuffle=True,
... )
>>> datamodule.setup()
>>> train_loader = datamodule.train_dataloader()
>>> val_loader = datamodule.val_dataloader()
>>> test_loader = datamodule.test_dataloader()
>>> predict_loader = datamodule.predict_dataloader()

Notes

• The setup method must be called before accessing dataloaders to ensure proper dataset splitting.

METHOD	DESCRIPTION
setup	Setup train and validation datasets with chosen splitting strategy.
predict_dataloader	Return dataloader for prediction on full dataset.
train_dataloader	Return dataloader for training dataset.
val_dataloader	Return dataloader for validation dataset, or None if no validation split.
test_dataloader	Return dataloader for testing on full dataset.
check_setup	Check if the datamodule has been set up and raise error if not.

Source code in spectre/data/datamodule.py

def __init__(
    self,
    dataset: DatasetWeightedType | dict,
    batch_size: int = 1000,
    train_stride: int = 1,
    val: float = 0.2,
    use_random_split: bool = True,
    shuffle: bool = False,
    stratify: bool = False,
) -> None:
    super().__init__()

    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(train_stride, int):
        raise TypeError("train_stride must be an integer.")
    check_in_interval(train_stride, "(0, inf)")
    self.train_stride = train_stride

    if not isinstance(val, (int, float)):
        raise TypeError("val must be a number.")
    check_in_interval(val, "[0, 1]")
    self.val = val

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

    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 dataset is None:
        raise TypeError("Dataset cannot be None.")
    elif isinstance(dataset, Subset):
        raise ValueError(
            "Subset datasets are not supported to avoid double "
            "indexing issues. Please pass the original dataset instead."
        )
    elif 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 TypeError(
            f"Unsupported dataset type: {type(dataset).__name__}. "
            f"Expected `DatasetWeighted`, `DatasetWeightedMemoryMapped`, "
            f"or `dict`."
        )

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

    self.dataset_strided = Subset(
        self.dataset,
        torch.arange(start=0, end=len(self.dataset), step=train_stride).tolist(),
    )

    self._train_dataset = None
    self._val_dataset = None
    self._test_dataset = None

    self._dataloaders = {}

    self._setup = False

Functions

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

Setup train and validation datasets with chosen splitting strategy.

Creates train and validation splits using either random or sequential splitting based on the validation fraction and splitting method specified during initialization.

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.py

def setup(self, stage: str | None = None) -> None:
    """
    Setup train and validation datasets with chosen splitting strategy.

    Creates train and validation splits using either random or sequential splitting
    based on the validation fraction and splitting method specified during
    initialization.

    Parameters
    ----------
    stage : str | None, optional, by default None
        PyTorch Lightning stage (unused but required by interface).
    """
    if self.val:
        n = len(self.dataset_strided)
        n_train = int((1.0 - self.val) * n)
        n_valid = n - n_train
        if self.use_random_split:
            self._train_dataset, self._val_dataset = random_split(
                self.dataset_strided, [n_train, n_valid]
            )
        else:
            train_indices = torch.arange(n_train).tolist()
            val_indices = torch.arange(n_train, n).tolist()
            self._train_dataset = Subset(self.dataset_strided, train_indices)
            self._val_dataset = Subset(self.dataset_strided, val_indices)
    else:
        self._train_dataset = self.dataset_strided

    self._setup = True

predict_dataloader() -> DataLoaderWeighted

Return dataloader for prediction on full dataset.

Source code in spectre/data/datamodule.py

def predict_dataloader(self) -> DataLoaderWeighted:
    """Return dataloader for prediction on full dataset."""
    self.check_setup()
    if "predict" not in self._dataloaders:
        self._dataloaders["predict"] = DataLoaderWeighted(
            dataset=self.dataset,
            batch_size=self.batch_size,
            shuffle=self.shuffle,
            stratify=self.stratify,
            drop_last=False,
        )
    return self._dataloaders["predict"]

train_dataloader() -> DataLoaderWeighted

Return dataloader for training dataset.

Source code in spectre/data/datamodule.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=self.shuffle,
            stratify=self.stratify,
            drop_last=True,
        )
    return self._dataloaders["train"]

val_dataloader() -> DataLoaderWeighted | None

Return dataloader for validation dataset, or None if no validation split.

Source code in spectre/data/datamodule.py

def val_dataloader(self) -> DataLoaderWeighted | None:
    """Return dataloader for validation dataset, or None if no validation split."""
    self.check_setup()
    if self.val:
        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=self.shuffle,
                stratify=self.stratify,
                drop_last=True,
            )
        return self._dataloaders["val"]
    else:
        return None

test_dataloader() -> DataLoaderWeighted

Return dataloader for testing on full dataset.

Source code in spectre/data/datamodule.py

def test_dataloader(self) -> DataLoaderWeighted:
    """Return dataloader for testing on full dataset."""
    self.check_setup()
    if "test" not in self._dataloaders:
        self._dataloaders["test"] = DataLoaderWeighted(
            dataset=self.dataset,
            batch_size=self.batch_size,
            shuffle=False,
            stratify=False,  # Don't stratify test set
            drop_last=False,
        )
    return self._dataloaders["test"]

check_setup() -> None

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

Source code in spectre/data/datamodule.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