mml.core.data_loading.lightning_datamodule

class MultiTaskDataModule

Bases: LightningDataModule

This class wraps one or multiple TaskDataset s for lightning. Given the respective TaskStruct s it takes care of setting up all correct data splits. It particularly interprets the following elements of the config:

• loaders: the ModalityLoader s

• preprocessing: the AugmentationModule

• augmentations: also AugmentationModule

• all aspects with respect to sampling and the dataloader

Importantly it provides the necessary lightning interface (e.g., setup(), train_dataloader(), etc.).

__init__(task_structs: List[TaskStruct], cfg: DictConfig, fold: int = 0)

get_cpu_transforms(struct: TaskStruct, phase: LearningPhase = LearningPhase.TRAIN) → AugmentationModule | AugmentationModuleContainer

Returns the necessary :param struct: :param phase: :return:

static get_dataset_balancing_weights(ds: TaskDataset) → Tensor

get_image_normalization(struct: TaskStruct) → Tuple[RGBInfo | None, RGBInfo | None]

Returns the applied / required image normalization information.

Returns:

tuple of means and stds for each of the channels, in case no normalization is applied returns None for both

Return type:

Tuple[Optional[RGBInfo], Optional[RGBInfo]]

get_loader_kwargs_from_cfg(task_name: str, phase: LearningPhase = LearningPhase.TRAIN) → Dict[str, Any]

get_modality_loaders() → Dict[Modality, ModalityLoader]

Creates ModalityLoader instances from the config.

on_after_batch_transfer(batch: Any, dataloader_idx: int) → Any

Enables gpu augmentations after batch has been transferred to device.

predict_dataloader(*args, **kwargs) → CombinedLoader

prepare_data(*args, **kwargs)

setup(stage: str) → None

Implements the lightning interface to prepare the datamodule. In particular sets up the TaskDataset s

Parameters:

stage

Returns:

teardown(stage: str | None = None) → None

test_dataloader(*args, **kwargs) → CombinedLoader

train_dataloader(*args, **kwargs) → CombinedLoader

val_dataloader(*args, **kwargs) → CombinedLoader