make DataModule compatible with Python dataclass

[awaelchli]

🚀 Feature

Support the following:

@dataclass
class MyDataModule(LightningDataModule):
    pass

Motivation

To reduce boilerplate code is at the core of philosophy in Lightning. It should be compatible with dataclasses.

Code sample

Here is an example. It currently does not work as we have some internal attributes that don't play well with the dataclass.

import os
from dataclasses import dataclass

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer, LightningDataModule


class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


@dataclass
class BoringDataModule(LightningDataModule):

    batch_size: int = 2

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


class BoringModel(LightningModule):

    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

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

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        weights_summary=None,
    )
    trainer.fit(model, datamodule=BoringDataModule())


if __name__ == '__main__':
    run()

Alternatives

#3792 introduces save_hyperparameters() for the datamodule. However, I believe the dataclass approach here is not in conflict with that because both could be useful at the same time.

[QueshAnmak]

Hi, I would like to work on this issue, could you please assign it to me?

[awaelchli]

Hey! Yes, you can take it if you want. I haven't really had the time to look why it is not working. I don't know how difficult it would be.
Give it a try. We appreciate the help!

[QueshAnmak]

Hey, Is there any slack or discord server I could join to ask my queries?

[awaelchli]

Yes, feel free to join the PyTorch Lightning slack: https://join.slack.com/t/pytorch-lightning/shared_invite/zt-pw5v393p-qRaDgEk24~EjiZNBpSQFgQ

[QueshAnmak]

I set init to False (which by default is set to True). This seems to fix the issue. I think the "@Property" decorator on the attributes might be the cause of the issue.

How should I proceed?

[awaelchli]

Oh yes, but actually we want dataclass to generate an init for us. Otherwise we don't get any great value of a dataclass here.
But actually, all we have to do is this:

    def __post_init__(self):
        super().__init__()

Because the init generated by dataclass does not call super(), and that's required by the datamodule.

[adosar]

@awaelchli @QueshAnmak For a custom datamodule should self.save_hyperparmeters called under __post_init__, as following?

@dataclass
class MyDataModule(LightningDataModule):
    x: int = 32
    ...
    
    def __post_init_(self):
          super().__init__()
          self.save_hyperparameters()

Currently, self.save_hyperparameters behaves differently depending on where it is called. I don't think this is the expected behavior since self.hparams should be populated with the init_args.

The following example demonstrates that with the @dataclass decorator, self.save_hyperparameters doesn't capture the init_args.

from dataclasses import dataclass
import lightning as L

class Foo(L.LightningDataModule):
    def __init__(self, x=1, y=2):
        super().__init__()
        self.y = 'hello'
        self.save_hyperparameters()

@dataclass
class Bar(L.LightningDataModule):
    x: int = 1
    y: int = 2

    def __post_init__(self):
        super().__init__()
        self.y = 'hello'
        self.save_hyperparameters()

foo = Foo()
print(foo.hparams)

bar = Bar()
print(bar.hparams)

# Output
"x": 1                                                                                                                                                                                        
"y": 2                                                                                                                                                                                        
"x": 1                                                                                                                                                                                        
"y": hello