Dataset Manager

class dipm.data.dataset_manager.DatasetManager(config: DatasetManagerConfig, dataset_creation: DatasetCreation)

Prepares a dataset by calculating statistics and storing masks. This is a part of the DatasetManager class.

Parameters:

config (DatasetManagerConfig) – Configuration object containing the dataset paths etc.

__init__(config: DatasetManagerConfig, dataset_creation: DatasetCreation)

prepare_datasets(dataset_info: DatasetInfo | None = None) → DatasetInfo

Calculate statistics and store masks. Shuffling and splitting is disabled.

Parameters:

dataset_info (DatasetInfo, optional) – The dataset information. If provided, only to remove unseen elements and empty graphs.

Returns:

The calculated/updated dataset information.

get_loaders(devices: list[Device] | None = None) → tuple[DataLoader, DataLoader | None, DataLoader | None]

Get data loaders. Shuffling and splitting is enabled.

Parameters:

devices (list[jax.Device], optional) – The devices to use for training. If provided, parallel training is enabled.

Returns:

A tuple of training, validation (if provided), and test (if provided) dataloaders.