DataModule

class dmme.data_modules.DataModule(batch_size)

LightningDataModule with defaults for generative modeling

> Defaults are set from DDPM.

setup_train and setup_test is used for preparing training and test sets. In practice, they both use training sets but augmentations are only applied on setup_train

Prepares DataLoader`s with good defaults with batch size set from `__init__.

Parameters:

batch_size (int) – batch size for DataLoader

setup_train()

Prepare training set

setup_test()

Prepare test set

setup(stage: str)

Prepare dataset for training or testing

train_dataloader()

DataLoader with good defaults

automatically sets num_workers based on cpu count.

test_dataloader()

DataLoader with good defaults

automatically sets num_workers based on cpu count.

dmme.data_modules.cpu_count(*args, **kwargs)

returns cpu count from multiprocessing package