Do LightningDataModule support multi train dataloader like LightningModule?

[kleinzcy]

Can I define multi training dataloaders in LightningDataModule?

    def train_dataloader(self):
        lab_loader = torch.utils.data.DataLoader(
                         self.train_subset_lab_train,
                         batch_size=self.batch_size,
                         shuffle=True,
                         num_workers=self.num_workers,
                         pin_memory=True,
                         drop_last=True,
                    )
        unlab_loader = torch.utils.data.DataLoader(
                         self.train_subset_unlab_train,
                         batch_size=self.batch_size*self.uratio,
                         shuffle=True,
                         num_workers=self.num_workers,
                         pin_memory=True,
                         drop_last=True,
                    )

        return [lab_loader, unlab_loader]

[tshu-w]

yes

[kleinzcy]

I have tried to implement it in the LightningDataModule, but I found it returns a dataloader object.

class A(LightingModule):
     ...
    def training_step(self, batch, _):
        data, labels = batch

I know it works when I define multi-dataloader in the LightingModule. But when I implement multi-dataloader in the LightningDataModule, and combine this with LightingModule, it doesn't work. I have not figured out why.

[tshu-w]

Both LightningModule and LightningDataModule integrate the class DataHook, they should be the same (in my case both support multiple dataloaders). You should provide sample code to reproduce your error

[kleinzcy]

The code below can reproduce my error:

import torch
import torchvision
import torchvision.transforms as T
import pytorch_lightning as pl
from torchvision.datasets import MNIST
from torchvision.datasets import CIFAR100,CIFAR10
import torch.nn as nn
import numpy as np
import os


class DataModule(pl.LightningDataModule):
    def __init__(self):
        super().__init__()
        self.transform = T.ToTensor()

    def setup(self, stage=None):
        self.mnist_data = MNIST('./data', train=True, download=True, transform=self.transform)
        print('mnist: ', len(self.mnist_data))
        self.cifar10_data = CIFAR10('./data', train=True, download=True, transform=self.transform)
        print('cifar: ', len(self.cifar10_data))

    def train_dataloader(self):
        A_loader = torch.utils.data.DataLoader(
                         self.mnist_data,
                         batch_size=2,
                         shuffle=True,
                         num_workers=0,
                         pin_memory=True,
                         drop_last=True,
                    )
        B_loader = torch.utils.data.DataLoader(
                         self.cifar10_data,
                         batch_size=2,
                         shuffle=True,
                         num_workers=0,
                         pin_memory=True,
                         drop_last=True,
                    )

        return [A_loader, B_loader]


class Sample(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = nn.Linear(10, 10)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, _):
        import ipdb
        ipdb.set_trace()
        # batch is a dataloader object
        views, labels = batch

        loss = 0
        return loss


def main():
    dm = DataModule()

    model = Sample()
    trainer = pl.Trainer()
    trainer.fit(model, dm)


if __name__ == "__main__":
    main()

[tchaton]

@kleinzcy @tshu-w,

PyTorch Lightning supports multiple data loaders but they will be sampled at the same time and return a batch composed of a batch sample for each dataloaders return for the train_dataloader method.

Alternatively, if you want to make this sequential, you can implement a wrapper which will sample from one and the other dataloader.

class SequentialLoader(Iterator):

    def __init__(self, *dataloaders):
        self.dataloaders = dataloaders

    def __len__(self):
        return sum([len(dl) for dl in self.dataloaders])

    def __iter__(self):
        for dl in self.dataloaders:
            dataloader_iter = iter(dl)
            for batch in dataloader_iter:
                yield batch

[jlehrer1]

What class is Iterator from? I'm facing precisely this problem now, wanting to sample sequentially from multiple dataloaders.

[tshu-w]

@jlehrer1 Hello, you can refer to my implementation:

from torch.utils.data import DataLoader


class SequentialLoader:
    def __init__(self, *dataloaders: DataLoader):
        self.dataloaders = dataloaders

    def __len__(self):
        return sum(len(d) for d in self.dataloaders)

    def __iter__(self):
        for dataloader in self.dataloaders:
            yield from dataloader

and for train_dataloader:

    def train_dataloader(self) -> TRAIN_DATALOADERS:
        dataloaders = (
            DataLoader(
                dataset=dataset,
                batch_size=self.hparams.batch_size,
                num_workers=self.hparams.num_workers,
                pin_memory=self.hparams.pin_memory,
                collate_fn=self.collate_fn,
                persistent_workers=self.hparams.num_workers > 0,
                shuffle=True,
            )
            for dataset in self.datasets
        )
        return SequentialLoader(*dataloaders)