Understanding PyTorch Lightning DataModules

PyTorch Lightning aims to make PyTorch code more structured and readable and that not just limited to the PyTorch Model but also the data itself. In PyTorch we use DataLoaders to train or test our model. While we can use DataLoaders in PyTorch Lightning to train the model too, PyTorch Lightning also provides us with a better approach called DataModules. DataModule is a reusable and shareable class that encapsulates the DataLoaders along with the steps required to process data. Creating dataloaders can get messy that’s why it’s better to club the dataset in the form of DataModule. Its recommended that you know how to define a neural network using PyTorch Lightning.

Installing PyTorch Lightning:

Installing Lightning is the same as that of any other library in python.

pip install pytorch-lightning

Or if you want to install it in a conda environment you can use the following command:-

conda install -c conda-forge pytorch-lightning

Pytorch Lightning DataModule Format

To define a Lightning DataModule we follow the following format:-

import pytorch-lightning as pl
from torch.utils.data import random_split, DataLoader

class DataModuleClass(pl.LightningDataModule):
    def __init__(self):
        #Define required parameters here
    
    def prepare_data(self):
        # Define steps that should be done
    # on only one GPU, like getting data.
    
    def setup(self, stage=None):
        # Define steps that should be done on 
    # every GPU, like splitting data, applying
    # transform etc.
    
    def train_dataloader(self):
        # Return DataLoader for Training Data here
    
    def val_dataloader(self):
        # Return DataLoader for Validation Data here
    
    def test_dataloader(self):
        # Return DataLoader for Testing Data here

Note: The names of the above functions should be exactly the same.

Understanding the DataModule Class

For this article, I’ll be using MNIST data as an example. As we can see, the first requirement to create a Lightning DataModule is to inherit the LightningDataModule class in pytorch-lightning:

import pytorch-lightning as pl
from torch.utils.data import random_split, DataLoader

class DataModuleMNIST(pl.LightningDataModule):

__init__() method:

It is used to store information regarding batch size, transforms, etc. 

def __init__(self):
    super().__init__()
    self.download_dir = ''
    self.batch_size = 32
    self.transform = transforms.Compose([
        transforms.ToTensor()
    ])

prepare_data() method:

This method is used to define the processes that are meant to be performed by only one GPU. It’s usually used to handle the task of downloading the data. 

def prepare_data(self):
    datasets.MNIST(self.download_dir,
           train=True, download=True)
           
    datasets.MNIST(self.download_dir, train=False,        
           download=True)

setup() method:

This method is used to define the process that is meant to be performed by all the available GPU. It’s usually used to handle the task of loading the data. 

def setup(self, stage=None):
    data = datasets.MNIST(self.download_dir,
             train=True, transform=self.transform)
             
    self.train_data, self.valid_data = random_split(data, [55000, 5000])
        
    self.test_data = datasets.MNIST(self.download_dir,
                        train=False, transform=self.transform)

train_dataloader() method:

This method is used to create a training data dataloader. In this function, you usually just return the dataloader of training data.

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

val_dataloader() method:

This method is used to create a validation data dataloader. In this function, you usually just return the dataloader of validation data.

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

test_dataloader() method:

This method is used to create a testing data dataloader. In this function, you usually just return the dataloader of testing data.

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

Training Pytorch Lightning Model Using DataModule:

In Pytorch Lighting, we use Trainer() to train our model and in this, we can pass the data as DataLoader or DataModule. Let’s use the model I defined in this article here as an example:

class model(pl.LightningModule): 
    def __init__(self): 
        super(model, self).__init__() 
        self.fc1 = nn.Linear(28*28, 256) 
        self.fc2 = nn.Linear(256, 128) 
        self.out = nn.Linear(128, 10) 
        self.lr = 0.01
        self.loss = nn.CrossEntropyLoss() 
  
    def forward(self, x): 
        batch_size, _, _, _ = x.size() 
        x = x.view(batch_size, -1) 
        x = F.relu(self.fc1(x)) 
        x = F.relu(self.fc2(x)) 
        return self.out(x) 
  
    def configure_optimizers(self): 
        return torch.optim.SGD(self.parameters(), lr=self.lr) 
  
    def training_step(self, train_batch, batch_idx): 
        x, y = train_batch 
        logits = self.forward(x) 
        loss = self.loss(logits, y) 
        return loss 
  
    def validation_step(self, valid_batch, batch_idx): 
        x, y = valid_batch 
        logits = self.forward(x) 
        loss = self.loss(logits, y)

Now to train this model we’ll create a Trainer() object and fit() it by passing our model and datamodules as parameters.

clf = model() 
mnist = DataModuleMNIST() 
trainer = pl.Trainer(gpus=1) 
trainer.fit(clf, mnist)

Below the full implementation:

Output: