LightningDataModule

class pytorch_lightning.core.LightningDataModule

Bases: pytorch_lightning.core.hooks.CheckpointHooks, pytorch_lightning.core.hooks.DataHooks, pytorch_lightning.core.mixins.hparams_mixin.HyperparametersMixin

A DataModule standardizes the training, val, test splits, data preparation and transforms. The main advantage is consistent data splits, data preparation and transforms across models.

Example:

class MyDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
    def prepare_data(self):
        # download, split, etc...
        # only called on 1 GPU/TPU in distributed
    def setup(self, stage):
        # make assignments here (val/train/test split)
        # called on every process in DDP
    def train_dataloader(self):
        train_split = Dataset(...)
        return DataLoader(train_split)
    def val_dataloader(self):
        val_split = Dataset(...)
        return DataLoader(val_split)
    def test_dataloader(self):
        test_split = Dataset(...)
        return DataLoader(test_split)
    def teardown(self):
        # clean up after fit or test
        # called on every process in DDP

classmethod add_argparse_args(parent_parser, **kwargs)

Extends existing argparse by default LightningDataModule attributes.

Example:

parser = ArgumentParser(add_help=False)
parser = LightningDataModule.add_argparse_args(parser)

Return type:

ArgumentParser

classmethod from_argparse_args(args, **kwargs)

Create an instance from CLI arguments.

Parameters:

• args (Union[Namespace, ArgumentParser]) – The parser or namespace to take arguments from. Only known arguments will be parsed and passed to the LightningDataModule.

• **kwargs – Additional keyword arguments that may override ones in the parser or namespace. These must be valid DataModule arguments.

Example:

module = LightningDataModule.from_argparse_args(args)

classmethod from_datasets(train_dataset=None, val_dataset=None, test_dataset=None, predict_dataset=None, batch_size=1, num_workers=0)

Create an instance from torch.utils.data.Dataset.

Parameters:

• train_dataset (Union[Dataset, Sequence[Dataset], Mapping[str, Dataset], None]) – (optional) Dataset to be used for train_dataloader()

• val_dataset (Union[Dataset, Sequence[Dataset], None]) – (optional) Dataset or list of Dataset to be used for val_dataloader()

• test_dataset (Union[Dataset, Sequence[Dataset], None]) – (optional) Dataset or list of Dataset to be used for test_dataloader()

• predict_dataset (Union[Dataset, Sequence[Dataset], None]) – (optional) Dataset or list of Dataset to be used for predict_dataloader()

• batch_size (int) – Batch size to use for each dataloader. Default is 1.

• num_workers (int) – Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. Number of CPUs available.

classmethod get_init_arguments_and_types()

Scans the DataModule signature and returns argument names, types and default values.

Returns:

(argument name, set with argument types, argument default value).

Return type:

List with tuples of 3 values

classmethod load_from_checkpoint(checkpoint_path, hparams_file=None, **kwargs)

Primary way of loading a datamodule from a checkpoint. When Lightning saves a checkpoint it stores the arguments passed to __init__ in the checkpoint under "datamodule_hyper_parameters".

Any arguments specified through **kwargs will override args stored in "datamodule_hyper_parameters".

Parameters:

• checkpoint_path (Union[str, Path, IO]) – Path to checkpoint. This can also be a URL, or file-like object

• hparams_file (Union[str, Path, None]) –

  Optional path to a .yaml or .csv file with hierarchical structure as in this example:

  dataloader:
      batch_size: 32
  

  You most likely won’t need this since Lightning will always save the hyperparameters to the checkpoint. However, if your checkpoint weights don’t have the hyperparameters saved, use this method to pass in a .yaml file with the hparams you’d like to use. These will be converted into a dict and passed into your LightningDataModule for use.

  If your datamodule’s hparams argument is Namespace and .yaml file has hierarchical structure, you need to refactor your datamodule to treat hparams as dict.

• **kwargs – Any extra keyword args needed to init the datamodule. Can also be used to override saved hyperparameter values.

Returns:

LightningDataModule instance with loaded weights and hyperparameters (if available).

Note

load_from_checkpoint is a class method. You should use your LightningDataModule class to call it instead of the LightningDataModule instance.

Example:

# load weights without mapping ...
datamodule = MyLightningDataModule.load_from_checkpoint('path/to/checkpoint.ckpt')

# or load weights and hyperparameters from separate files.
datamodule = MyLightningDataModule.load_from_checkpoint(
    'path/to/checkpoint.ckpt',
    hparams_file='/path/to/hparams_file.yaml'
)

# override some of the params with new values
datamodule = MyLightningDataModule.load_from_checkpoint(
    PATH,
    batch_size=32,
    num_workers=10,
)

load_state_dict(state_dict)

Called when loading a checkpoint, implement to reload datamodule state given datamodule state_dict.

Parameters:

state_dict (Dict[str, Any]) – the datamodule state returned by state_dict.

Return type:

None

state_dict()

Called when saving a checkpoint, implement to generate and save datamodule state.

Return type:

Dict[str, Any]

Returns:

A dictionary containing datamodule state.