Inheritance of Data Loaders

[rusty1s]

🐛 Bug

This is the official bug report of my comment earlier (see #10680 (comment)).

import pytorch_lightning as pl
import torch


class Model(pl.LightningModule):
    def __init__(self, in_channels):
        super().__init__()
        self.lin = torch.nn.Linear(in_channels, 1)

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

    def training_step(self, x, batch_idx):
        loss = self(x).mean()
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)


class MyBaseDataLoader(torch.utils.data.DataLoader):
    pass


class MyDataLoader(MyBaseDataLoader):
    def __init__(self, data: torch.Tensor, *args, **kwargs):
        self.data = data
        super().__init__(range(data.size(0)), *args, **kwargs,
                         collate_fn=self.sample_fn)

    def sample_fn(self, indices):
        assert isinstance(self.data, torch.Tensor)
        index = torch.tensor(indices)
        return self.data[index]


class DataModule(pl.LightningDataModule):
    def __init__(self, data: torch.Tensor):
        super().__init__()
        self.data = data

    def train_dataloader(self):
        return MyDataLoader(self.data, batch_size=32)


model = Model(in_channels=32)
datamodule = DataModule(torch.randn(100, 32))

trainer = pl.Trainer(accelerator='gpu', devices=1, max_epochs=10)
trainer.fit(model, datamodule)

which crashes with

  File "/home/matthias/miniconda3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 530, in __next__
    data = self._next_data()
  File "/home/matthias/miniconda3/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 570, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/matthias/miniconda3/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    return self.collate_fn(data)
  File "/home/matthias/github/pytorch_geometric/alpha/PL.py", line 32, in sample_fn
    assert isinstance(self.data, torch.Tensor)
AssertionError

Notably, the code runs through for pytorch-lightning==1.5.10 and when I let MyDataLoader inherit directly from torch.utils.data.DataLoader.

I personally think this issue is related to overriding both __init__ from MyBaseDataLoader and MyDataLoader here.

• PyTorch Lightning Version (e.g., 1.5.0): 1.6.0
• PyTorch Version (e.g., 1.10): 1.10
• Python version (e.g., 3.9): 3.9
• OS (e.g., Linux): Linux
• CUDA/cuDNN version: 11.3

cc @justusschock @awaelchli @ninginthecloud