DataModule#

class DataModule(problem, train_size, val_size, test_size, batch_size, batching_mode, automatic_batching, shuffle, num_workers, pin_memory)[source]#

Bases: LightningDataModule

An extension of the Lightning data module for managing PINA condition datasets.

The data module handles train/validation/test dataset splitting, condition subset creation, dataloader construction, and batching coordination across multiple conditions.

Dataset splitting is performed independently for each condition, and the resulting subsets are wrapped into _ConditionSubset objects. Dataloaders are then created and aggregated according to the selected batching strategy.

Example:
>>> import torch
>>> from pina import LabelTensor
>>> from pina.condition import Condition
>>> from pina.problem import BaseProblem
>>> class MyProblem(BaseProblem):
...     def __init__(self):
...         super().__init__()
...         pts = LabelTensor(torch.randn(100, 2), labels=["x", "y"])
...         self.conditions = {"cond1": Condition(input=pts)}
>>> problem = MyProblem()
>>> dm = DataModule(problem, train_size=0.8, val_size=0.1,
...     test_size=0.1, batch_size=32, batching_mode="common_batch_size",
...     automatic_batching=False, shuffle=True, num_workers=0,
...     pin_memory=False)
>>> dm.setup("fit")
>>> list(dm.train_datasets.keys())
['cond1']

Initialization of the DataModule class.

Parameters:

  • problem (BaseProblem) – The problem containing the conditions and sampled data used to construct datasets and dataloaders.

  • train_size (float) – The fraction of samples assigned to the training split. Must belong to the interval [0, 1].

  • val_size (float) – The fraction of samples assigned to the validation split. Must belong to the interval [0, 1].

  • test_size (float) – The fraction of samples assigned to the test split. Must belong to the interval [0, 1].

  • batch_size (int) – The number of samples per batch. If None, the entire dataset is processed as a single batch.

  • batching_mode (str) – The strategy used to aggregate batches across dataloaders. Available options are "common_batch_size" for uniform batch sizes across conditions, "proportional" for batch sizes proportional to dataset sizes, and "separate_conditions" for iterating through each condition separately.

  • automatic_batching (bool) – Whether PyTorch automatic batching should be enabled. If True, dataset elements are retrieved individually and collated into batches by the dataloader. If False, entire subsets are retrieved directly from the condition object.

  • shuffle (bool) – Whether condition samples should be shuffled before splitting.

  • num_workers (int) – The number of worker processes used by dataloaders.

  • pin_memory (bool) – Whether pinned memory should be enabled during data loading.

Raises:

  • UserWarning – If num_workers is set to non-default value while batch_size is None.

  • UserWarning – If pin_memory is set to True while batch_size is None.

setup(stage=None)[source]#

Create dataset subsets for the requested execution stage.

Depending on the selected stage, it initializes the train_datasets, the val_datasets, or the test_datasets attributes. Each dataset is represented as a mapping between condition names and _ConditionSubset instances.

Parameters:

stage (str) – The execution stage. Available options are "fit" for training/validation and "test" for testing. If None, both training/validation and testing datasets are created. Default is None.

Raises:

ValueError – If the provided stage is invalid.

transfer_batch_to_device(batch, device, _)[source]#

Transfer a batch to the target device.

The method transfers all condition batches contained in the aggregated batch dictionary to the specified device.

Parameters:

  • batch (dict) – The mapping between the condition names and the condition batches.

  • device (torch.device) – The target device.

  • _ – Placeholder argument, not used.

Returns:

A list of tuples containing condition names and transferred batches.

Return type:

list[tuple[str, Any]]

train_dataloader()[source]#

Create the aggregated train dataloader.

Returns:

The aggregated dataloader coordinating all train condition dataloaders.

Return type:

_Aggregator

val_dataloader()[source]#

Create the aggregated validation dataloader.

Returns:

The aggregated dataloader coordinating all validation condition dataloaders.

Return type:

_Aggregator

test_dataloader()[source]#

Create the aggregated test dataloader.

Returns:

The aggregated dataloader coordinating all test condition dataloaders.

Return type:

_Aggregator