---
title: LightningDataModule
description: >-
  A datamodule is a shareable, reusable class that encapsulates all the steps
  needed to process data:
canonical: 'https://lightning.ai/docs/pytorch/stable/reference/data/datamodule'
---
# LightningDataModule {#lightningdatamodule}

A datamodule is a shareable, reusable class that encapsulates all the steps needed to process data:

[Video: https://pl-public-data.s3.amazonaws.com/assets_lightning/pt_dm_vid.mp4](https://pl-public-data.s3.amazonaws.com/assets_lightning/pt_dm_vid.mp4)

A datamodule encapsulates the five steps involved in data processing in PyTorch:

1. Download / tokenize / process.
2. Clean and (maybe) save to disk.
3. Load inside `Dataset`.
4. Apply transforms (rotate, tokenize, etc...).
5. Wrap inside a `DataLoader`.

This class can then be shared and used anywhere:

```python
model = LitClassifier()
trainer = Trainer()

imagenet = ImagenetDataModule()
trainer.fit(model, datamodule=imagenet)

cifar10 = CIFAR10DataModule()
trainer.fit(model, datamodule=cifar10)
```

## Why do I need a DataModule? {#why-do-i-need-a-datamodule}

In normal PyTorch code, the data cleaning/preparation is usually scattered across many files. This makes
sharing and reusing the exact splits and transforms across projects impossible.

Datamodules are for you if you ever asked the questions:

- what splits did you use?
- what transforms did you use?
- what normalization did you use?
- how did you prepare/tokenize the data?

## What is a DataModule? {#what-is-a-datamodule}

The `LightningDataModule`  is a convenient way to manage data in PyTorch Lightning.
It encapsulates training, validation, testing, and prediction dataloaders, as well as any necessary steps for data processing,
downloads, and transformations. By using a `LightningDataModule`, you can
easily develop dataset-agnostic models, hot-swap different datasets, and share data splits and transformations across projects.

Here's a simple PyTorch example:

```python
# regular PyTorch
test_data = MNIST(my_path, train=False, download=True)
predict_data = MNIST(my_path, train=False, download=True)
train_data = MNIST(my_path, train=True, download=True)
train_data, val_data = random_split(train_data, [55000, 5000])

train_loader = DataLoader(train_data, batch_size=32)
val_loader = DataLoader(val_data, batch_size=32)
test_loader = DataLoader(test_data, batch_size=32)
predict_loader = DataLoader(predict_data, batch_size=32)
```

The equivalent DataModule just organizes the same exact code, but makes it reusable across projects.

```python
class MNISTDataModule(L.LightningDataModule):
    def __init__(self, data_dir: str = "path/to/dir", batch_size: int = 32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def setup(self, stage: str):
        self.mnist_test = MNIST(self.data_dir, train=False)
        self.mnist_predict = MNIST(self.data_dir, train=False)
        mnist_full = MNIST(self.data_dir, train=True)
        self.mnist_train, self.mnist_val = random_split(
            mnist_full, [55000, 5000], generator=torch.Generator().manual_seed(42)
        )

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=self.batch_size)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=self.batch_size)

    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, batch_size=self.batch_size)

    def teardown(self, stage: str):
        # Used to clean-up when the run is finished
        ...
```

But now, as the complexity of your processing grows (transforms, multiple-GPU training), you can
let Lightning handle those details for you while making this dataset reusable so you can share with
colleagues or use in different projects.

```python
mnist = MNISTDataModule(my_path)
model = LitClassifier()

trainer = Trainer()
trainer.fit(model, mnist)
```

Here's a more realistic, complex DataModule that shows how much more reusable the datamodule is.

```python
import lightning as L
from torch.utils.data import random_split, DataLoader

# Note - you must have torchvision installed for this example
from torchvision.datasets import MNIST
from torchvision import transforms

class MNISTDataModule(L.LightningDataModule):
    def __init__(self, data_dir: str = "./"):
        super().__init__()
        self.data_dir = data_dir
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage: str):
        # Assign train/val datasets for use in dataloaders
        if stage == "fit":
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(
                mnist_full, [55000, 5000], generator=torch.Generator().manual_seed(42)
            )

        # Assign test dataset for use in dataloader(s)
        if stage == "test":
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

        if stage == "predict":
            self.mnist_predict = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=32)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=32)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=32)

    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, batch_size=32)
```

## LightningDataModule API {#lightningdatamodule-api}

To define a DataModule the following methods are used to create train/val/test/predict dataloaders:

- [prepare_data](/docs/pytorch/2.6.6/reference/data/datamodule.md#prepare-data) (how to download, tokenize, etc...)
- [setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup) (how to split, define dataset, etc...)
- [train_dataloader](/docs/pytorch/2.6.6/reference/data/datamodule.md#train-dataloader)
- [val_dataloader](/docs/pytorch/2.6.6/reference/data/datamodule.md#val-dataloader)
- [test_dataloader](/docs/pytorch/2.6.6/reference/data/datamodule.md#test-dataloader)
- [predict_dataloader](/docs/pytorch/2.6.6/reference/data/datamodule.md#predict-dataloader)

### prepare_data {#prepare-data}

Downloading and saving data with multiple processes (distributed settings) will result in corrupted data. Lightning
ensures the `prepare_data` is called only within a single process on CPU,
so you can safely add your downloading logic within. In case of multi-node training, the execution of this hook
depends upon [prepare_data_per_node](/docs/pytorch/2.6.6/reference/data/datamodule.md#prepare-data-per-node). `setup` is called after
`prepare_data` and there is a barrier in between which ensures that all the processes proceed to `setup` once the data is prepared and available for use.

- download, i.e. download data only once on the disk from a single process
- tokenize. Since it's a one time process, it is not recommended to do it on all processes
- etc...

```python
class MNISTDataModule(L.LightningDataModule):
    def prepare_data(self):
        # download
        MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor())
        MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())
```

> **Warning**
>
> `prepare_data` is called from the main process. It is not recommended to assign state here (e.g. `self.x = y`) since it is called on a single process and if you assign
> states here then they won't be available for other processes.

### setup {#setup}

There are also data operations you might want to perform on every GPU. Use `setup` to do things like:

- count number of classes
- build vocabulary
- perform train/val/test splits
- create datasets
- apply transforms (defined explicitly in your datamodule)
- etc...

```python
import lightning as L

class MNISTDataModule(L.LightningDataModule):
    def setup(self, stage: str):
        # Assign Train/val split(s) for use in Dataloaders
        if stage == "fit":
            mnist_full = MNIST(self.data_dir, train=True, download=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(
                mnist_full, [55000, 5000], generator=torch.Generator().manual_seed(42)
            )

        # Assign Test split(s) for use in Dataloaders
        if stage == "test":
            self.mnist_test = MNIST(self.data_dir, train=False, download=True, transform=self.transform)
```

For eg., if you are working with NLP task where you need to tokenize the text and use it, then you can do something like as follows:

```python
class LitDataModule(L.LightningDataModule):
    def prepare_data(self):
        dataset = load_Dataset(...)
        train_dataset = ...
        val_dataset = ...
        # tokenize
        # save it to disk

    def setup(self, stage):
        # load it back here
        dataset = load_dataset_from_disk(...)
```

This method expects a `stage` argument.
It is used to separate setup logic for `trainer.{fit,validate,test,predict}`.

> **Note**
>
> [setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup) is called from every process across all the nodes. Setting state here is recommended.

> **Note**
>
> [teardown](/docs/pytorch/2.6.6/reference/data/datamodule.md#teardown) can be used to clean up the state. It is also called from every process across all the nodes.

### train_dataloader {#train-dataloader}

Use the `train_dataloader` method to generate the training dataloader(s).
Usually you just wrap the dataset you defined in [setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup). This is the dataloader that the Trainer
`fit` method uses.

```python
import lightning as L

class MNISTDataModule(L.LightningDataModule):
    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=64)
```

### val_dataloader {#val-dataloader}

Use the `val_dataloader` method to generate the validation dataloader(s).
Usually you just wrap the dataset you defined in [setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup). This is the dataloader that the Trainer
`fit` and `validate` methods uses.

```python
import lightning as L

class MNISTDataModule(L.LightningDataModule):
    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=64)
```

### test_dataloader {#test-dataloader}

Use the `test_dataloader` method to generate the test dataloader(s).
Usually you just wrap the dataset you defined in [setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup). This is the dataloader that the Trainer
`test` method uses.

```python
import lightning as L

class MNISTDataModule(L.LightningDataModule):
    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=64)
```

### predict_dataloader {#predict-dataloader}

Use the `predict_dataloader` method to generate the prediction dataloader(s).
Usually you just wrap the dataset you defined in [setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup). This is the dataloader that the Trainer
`predict` method uses.

```python
import lightning as L

class MNISTDataModule(L.LightningDataModule):
    def predict_dataloader(self):
        return DataLoader(self.mnist_predict, batch_size=64)
```

### transfer_batch_to_device {#transfer-batch-to-device}

``lightning.pytorch.core.datamodule.LightningDataModule.transfer_batch_to_device``

### on_before_batch_transfer {#on-before-batch-transfer}

``lightning.pytorch.core.datamodule.LightningDataModule.on_before_batch_transfer``

### on_after_batch_transfer {#on-after-batch-transfer}

``lightning.pytorch.core.datamodule.LightningDataModule.on_after_batch_transfer``

### load_state_dict {#load-state-dict}

``lightning.pytorch.core.datamodule.LightningDataModule.load_state_dict``

### state_dict {#state-dict}

``lightning.pytorch.core.datamodule.LightningDataModule.state_dict``

### teardown {#teardown}

``lightning.pytorch.core.datamodule.LightningDataModule.teardown``

### prepare_data_per_node {#prepare-data-per-node}

If set to `True` will call `prepare_data()` on LOCAL_RANK=0 for every node.
If set to `False` will only call from NODE_RANK=0, LOCAL_RANK=0.

```python
class LitDataModule(LightningDataModule):
    def __init__(self):
        super().__init__()
        self.prepare_data_per_node = True
```

## Using a DataModule {#using-a-datamodule}

The recommended way to use a DataModule is simply:

```python
dm = MNISTDataModule()
model = Model()
trainer.fit(model, datamodule=dm)
trainer.test(datamodule=dm)
trainer.validate(datamodule=dm)
trainer.predict(datamodule=dm)
```

If you need information from the dataset to build your model, then run
[prepare_data](/docs/pytorch/2.6.6/reference/data/datamodule.md#prepare-data) and
[setup](/docs/pytorch/2.6.6/reference/data/datamodule.md#setup) manually (Lightning ensures
the method runs on the correct devices).

```python
dm = MNISTDataModule()
dm.prepare_data()
dm.setup(stage="fit")

model = Model(num_classes=dm.num_classes, width=dm.width, vocab=dm.vocab)
trainer.fit(model, dm)

dm.setup(stage="test")
trainer.test(datamodule=dm)
```

You can access the current used datamodule of a trainer via `trainer.datamodule` and the current used
dataloaders via the trainer properties `train_dataloader`,
`val_dataloaders`,
`test_dataloaders`, and
`predict_dataloaders`.

## DataModules without Lightning {#datamodules-without-lightning}

You can of course use DataModules in plain PyTorch code as well.

```python
# download, etc...
dm = MNISTDataModule()
dm.prepare_data()

# splits/transforms
dm.setup(stage="fit")

# use data
for batch in dm.train_dataloader():
    ...

for batch in dm.val_dataloader():
    ...

dm.teardown(stage="fit")

# lazy load test data
dm.setup(stage="test")
for batch in dm.test_dataloader():
    ...

dm.teardown(stage="test")
```

But overall, DataModules encourage reproducibility by allowing all details of a dataset to be specified in a unified
structure.

## Hyperparameters in DataModules {#hyperparameters-in-datamodules}

Like LightningModules, DataModules support hyperparameters with the same API.

```python
import lightning as L

class CustomDataModule(L.LightningDataModule):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.save_hyperparameters()

    def configure_optimizers(self):
        # access the saved hyperparameters
        opt = optim.Adam(self.parameters(), lr=self.hparams.lr)
```

Refer to `save_hyperparameters` in [lightning module](/docs/pytorch/2.6.6/core-api/lightning_module.md) for more details.

### Save DataModule state {#save-datamodule-state}

When a checkpoint is created, it asks every DataModule for their state. If your DataModule defines the *state_dict* and *load_state_dict* methods, the checkpoint will automatically track and restore your DataModules.

```python
import lightning as L

class LitDataModule(L.LightningDataModule):
    def state_dict(self):
        # track whatever you want here
        state = {"current_train_batch_index": self.current_train_batch_index}
        return state

    def load_state_dict(self, state_dict):
        # restore the state based on what you tracked in (def state_dict)
        self.current_train_batch_index = state_dict["current_train_batch_index"]
```

---

**Navigation**

- Parent: [Data](/docs/pytorch/2.6.6/reference/data.md)