base

BaseDataModule + its datasets/config classes.

class cneuromax.fitting.deeplearning.datamodule.base.Datasets(train=None, val=None, test=None, predict=None)

Bases: object

Holds phase-specific torch.utils.data.Dataset objects.

Using the word phase to not overload Lightning stage terminology used for fit, validate and test.

Parameters:

• train (Union[Dataset[Tensor | dict[str, Tensor]], Dataset, None], default: None)

• val (Union[Dataset[Tensor | dict[str, Tensor]], Dataset, None], default: None)

• test (Union[Dataset[Tensor | dict[str, Tensor]], Dataset, None], default: None)

• predict (Union[Dataset[Tensor | dict[str, Tensor]], Dataset, None], default: None)

class cneuromax.fitting.deeplearning.datamodule.base.BaseDataModuleConfig(data_dir='${config.data_dir}', device='${config.device}', max_per_device_batch_size=None, fixed_per_device_batch_size=None, fixed_per_device_num_workers=None)

Bases: object

Holds BaseDataModule config values.

Parameters:

• data_dir (str, default: '${config.data_dir}') – See data_dir.

• device (str, default: '${config.device}') – See device.

• max_per_device_batch_size (Optional[int], default: None) – See per_device_batch_size. Sets an upper bound on the aforementioned attribute.

• fixed_per_device_batch_size (Optional[int], default: None) – See per_device_batch_size. Setting this value skips the batch size search in find_good_per_device_batch_size() which is not recommended for resource efficiency.

• fixed_per_device_num_workers (Optional[int], default: None) – See per_device_num_workers. Setting this value skips the num workers search in find_good_per_device_num_workers() which is not recommended for resource efficiency.

class cneuromax.fitting.deeplearning.datamodule.base.BaseDataModule(config)

Bases: LightningDataModule, ABC

Base lightning.pytorch.core.LightningDataModule.

With <phase> being any of train, val, test or predict, subclasses need to properly define the datasets.<phase> attribute(s) for each desired phase.

Parameters:

config (BaseDataModuleConfig)

config

Type:

BaseDataModuleConfig

datasets

Type:

Datasets

collate_fn

See collate_fn argument in torch.utils.data.DataLoader.

Type:

Callable

pin_memory

Whether to copy tensors into device pinned memory before returning them (is set to True by default if device is "gpu").

Type:

bool

per_device_batch_size

Per-device number of samples to load per iteration. Temporary value (1) is overwritten in set_batch_size_and_num_workers().

Type:

int

per_device_num_workers

Per-device number of CPU processes to use for data loading (0 means that the data will be loaded by each device’s assigned CPU process). Temporary value (0) is later overwritten in set_batch_size_and_num_workers().

Type:

int

final load_state_dict(state_dict)

Replace instance attrib vals w/ state_dict vals.

Parameters:

state_dict (dict[str, int]) – Dictionary containing values to override per_device_batch_size & per_device_num_workers.

Return type:

None

final state_dict()

Returns instance attribute values.

Return type:

dict[str, int]

Returns:

A new dictionary containing attribute values

per_device_batch_size & per_device_num_workers.

final x_dataloader(dataset, *, shuffle=True)

Generic torch.utils.data.DataLoader factory method.

Parameters:

• dataset (Union[Dataset[Tensor], Dataset, None]) – A torch Dataset to wrap with a torch.utils.data.DataLoader

• shuffle (bool, default: True) – Whether to shuffle the dataset when iterating over it.

Raises:

AttributeError – If dataset is None.

Return type:

DataLoader[Tensor]

Returns:

A new

torch.utils.data.DataLoader instance wrapping the dataset argument.

final train_dataloader()

Calls x_dataloader() w/ datasets .train.

Return type:

DataLoader[Tensor]

Returns:

A new training

torch.utils.data.DataLoader instance.

final val_dataloader()

Calls x_dataloader() w/ datasets .val.

Return type:

DataLoader[Tensor]

Returns:

A new validation

torch.utils.data.DataLoader instance.

final test_dataloader()

Calls x_dataloader() w/ datasets .test.

Return type:

DataLoader[Tensor]

Returns:

A new testing

torch.utils.data.DataLoader instance.

final predict_dataloader()

Calls x_dataloader() w/ datasets .predict.

Return type:

DataLoader[Tensor]

Returns:

A new prediction

torch.utils.data.DataLoader instance that does not shuffle the dataset.