Improve Design of DataModule

[awaelchli]

🚀 Feature

Motivation

Since the introduction of datamodules we have duplicated code. For example, the train/val/test_dataloader methods in LightningModule and DataModule have the same name, same docs, same signature but live in two files.
Also the argparse methods are copy pasted from Trainer, it will become impossible to maintain this in the future.

Pitch

Factor out the hook definitions:

class DataHooks:

    def prepare_data(...):
    """ docs... """

    def setup(...):
    """ docs... """

    def train_dataloader(...):
    """ docs... """

    etc.

and then in the other classes, inherit from it:

class LightningModule(..., ModelHooks, DataHooks, ...)
    ...

class DataModule(..., DataHooks, ...)
    ...

Note also:

• we even have duplicated tests

Alternatives

I see no alternatives. This is the best way I know.

@PyTorchLightning/core-contributors makes sense?

[Borda]

yes, just make them abstract...

[williamFalcon]

A few things to keep in mind:

1. We want datamodules to be usable even if you're not using lightning for training.
2. We are moving away from Mixins.

[awaelchli]

We want datamodules to be usable even if you're not using lightning for training.

This won't be changed by my suggestion. It's just refactoring, nothing else.

We are moving away from Mixins.

Which mixins? I am aware that in Lightning some classes are named mixins but conceptually they are actually not and should not be called that. Mixins have their place and one should not confuse multiple inheritance with mixins. Could you clarify what you mean?