mml.core.data_loading.lightning_datamodule

class MultiTaskDataModule

Bases: LightningDataModule

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

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

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

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

test_dataloader(*args, **kwargs) → CombinedLoader

train_dataloader(*args, **kwargs) → CombinedLoader

val_dataloader(*args, **kwargs) → CombinedLoader