---
title: LightningDataModule
description: >-
  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.
canonical: >-
  https://lightning.ai/docs/pytorch/stable/optional-api/api/lightning.pytorch.core.LightningDataModule
---
# LightningDataModule {#lightningdatamodule}

### `class lightning.pytorch.core.LightningDataModule` {#lightning.pytorch.core.LightningDataModule}

Bases: [`DataHooks`](/docs/pytorch/2.5.5/optional-api/api/lightning.pytorch.core.hooks.DataHooks.md#lightning.pytorch.core.hooks.DataHooks "lightning.pytorch.core.hooks.DataHooks"), [`HyperparametersMixin`](/docs/pytorch/2.5.5/optional-api/api/lightning.pytorch.core.mixins.HyperparametersMixin.md#lightning.pytorch.core.mixins.HyperparametersMixin "lightning.pytorch.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:

```python
import lightning as L
import torch.utils.data as data
from lightning.pytorch.demos.boring_classes import RandomDataset

class MyDataModule(L.LightningDataModule):
    def prepare_data(self):
        # download, IO, etc. Useful with shared filesystems
        # 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
        dataset = RandomDataset(1, 100)
        self.train, self.val, self.test = data.random_split(
            dataset, [80, 10, 10], generator=torch.Generator().manual_seed(42)
        )

    def train_dataloader(self):
        return data.DataLoader(self.train)

    def val_dataloader(self):
        return data.DataLoader(self.val)

    def test_dataloader(self):
        return data.DataLoader(self.test)

    def on_exception(self, exception):
        # clean up state after the trainer faced an exception
        ...

    def teardown(self):
        # clean up state after the trainer stops, delete files...
        # called on every process in DDP
        ...
```

#### `prepare_data_per_node` {#lightning.pytorch.core.LightningDataModule.prepare_data_per_node}

If True, each LOCAL\_RANK=0 will call prepare data. Otherwise only NODE\_RANK=0, LOCAL\_RANK=0 will prepare data.

#### `allow_zero_length_dataloader_with_multiple_devices` {#lightning.pytorch.core.LightningDataModule.allow_zero_length_dataloader_with_multiple_devices}

If True, dataloader with zero length within local rank is allowed. Default value is False.

#### `classmethod from_datasets(train_dataset=None, val_dataset=None, test_dataset=None, predict_dataset=None, batch_size=1, num_workers=0, **datamodule_kwargs)` {#lightning.pytorch.core.LightningDataModule.from_datasets}

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

Parameters:

*   **train\_dataset** (`Union`\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Optional dataset or iterable of datasets to be used for train\_dataloader()
    
*   **val\_dataset** (`Union`\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Optional dataset or iterable of datasets to be used for val\_dataloader()
    
*   **test\_dataset** (`Union`\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Optional dataset or iterable of datasets to be used for test\_dataloader()
    
*   **predict\_dataset** (`Union`\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)"), [`Iterable`](https://docs.python.org/3/library/collections.abc.html#collections.abc.Iterable "(in Python v3.15)")\[[`Dataset`](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Dataset "(in PyTorch v2.15)")\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – Optional dataset or iterable of datasets to be used for predict\_dataloader()
    
*   **batch\_size** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Batch size to use for each dataloader. Default is 1. This parameter gets forwarded to the `__init__` if the datamodule has such a name defined in its signature.
    
*   **num\_workers** ([`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")) – Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. Number of CPUs available. This parameter gets forwarded to the `__init__` if the datamodule has such a name defined in its signature.
    
*   **\*\*datamodule\_kwargs** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – Additional parameters that get passed down to the datamodule’s `__init__`.
    

Return type:

[`LightningDataModule`](#lightning.pytorch.core.LightningDataModule "lightning.pytorch.core.datamodule.LightningDataModule")

#### `load_from_checkpoint(checkpoint_path, map_location=None, hparams_file=None, **kwargs)` {#lightning.pytorch.core.LightningDataModule.load_from_checkpoint}

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`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Path`](https://docs.python.org/3/library/pathlib.html#pathlib.Path "(in Python v3.15)"), [`IO`](https://docs.python.org/3/library/typing.html#typing.IO "(in Python v3.15)")\]) – Path to checkpoint. This can also be a URL, or file-like object
    
*   **map\_location** (`Union`\[[`device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)"), [`Callable`](https://docs.python.org/3/library/typing.html#typing.Callable "(in Python v3.15)")\[\[[`UntypedStorage`](https://docs.pytorch.org/docs/stable/storage.html#torch.UntypedStorage "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)")\], [`Optional`](https://docs.python.org/3/library/typing.html#typing.Optional "(in Python v3.15)")\[[`UntypedStorage`](https://docs.pytorch.org/docs/stable/storage.html#torch.UntypedStorage "(in PyTorch v2.15)")\]\], [`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[`Union`\[[`device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\], `Union`\[[`device`](https://docs.pytorch.org/docs/stable/tensor_attributes.html#torch.device "(in PyTorch v2.15)"), [`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`int`](https://docs.python.org/3/builtins/functions.html#int "(in Python v3.15)")\]\], [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) – If your checkpoint saved a GPU model and you now load on CPUs or a different number of GPUs, use this to map to the new setup. The behaviour is the same as in [`torch.load()`](https://docs.pytorch.org/docs/stable/generated/torch.load.html#torch.load "(in PyTorch v2.15)").
    
*   **hparams\_file** (`Union`\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Path`](https://docs.python.org/3/library/pathlib.html#pathlib.Path "(in Python v3.15)"), [`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")\]) –
    
    Optional path to a `.yaml` or `.csv` file with hierarchical structure as in this example:
    
    ```python
    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`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)") and passed into your [`LightningDataModule`](#lightning.pytorch.core.LightningDataModule "lightning.pytorch.core.LightningDataModule") for use.
    
    If your datamodule’s `hparams` argument is [`Namespace`](https://docs.python.org/3/library/argparse.html#argparse.Namespace "(in Python v3.15)") and `.yaml` file has hierarchical structure, you need to refactor your datamodule to treat `hparams` as [`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)").
    
*   **\*\*kwargs** ([`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")) – Any extra keyword args needed to init the datamodule. Can also be used to override saved hyperparameter values.
    

Return type:

`Self`

Returns:

[`LightningDataModule`](#lightning.pytorch.core.LightningDataModule "lightning.pytorch.core.LightningDataModule") instance with loaded weights and hyperparameters (if available).

Note

`load_from_checkpoint` is a **class** method. You must use your [`LightningDataModule`](#lightning.pytorch.core.LightningDataModule "lightning.pytorch.core.LightningDataModule") **class** to call it instead of the [`LightningDataModule`](#lightning.pytorch.core.LightningDataModule "lightning.pytorch.core.LightningDataModule") instance, or a `TypeError` will be raised.

Example:

```python
# 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)` {#lightning.pytorch.core.LightningDataModule.load_state_dict}

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

Parameters:

**state\_dict** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]) – the datamodule state returned by `state_dict`.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `on_exception(exception)` {#lightning.pytorch.core.LightningDataModule.on_exception}

Called when the trainer execution is interrupted by an exception.

Return type:

[`None`](https://docs.python.org/3/builtins/constants.html#None "(in Python v3.15)")

#### `state_dict()` {#lightning.pytorch.core.LightningDataModule.state_dict}

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

Return type:

[`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict "(in Python v3.15)")\[[`str`](https://docs.python.org/3/builtins/stdtypes.html#str "(in Python v3.15)"), [`Any`](https://docs.python.org/3/library/typing.html#typing.Any "(in Python v3.15)")\]

Returns:

A dictionary containing datamodule state.

---

**Navigation**

- Parent: [Api](/docs/pytorch/2.5.5/optional-api/api.md)