Avoid patching LightningModule methods during training

[awaelchli]

🚀 Feature

Can we implement the dataloaders without 🐒-patching the methods in LightningModule?

Motivation

Currently, we patch the LightningModule methods in the trainer when also a DataModule is used.
https://github.com/PyTorchLightning/pytorch-lightning/blob/5157ba55095a6a9f93ec1976aac877c87b00158f/pytorch_lightning/trainer/connectors/data_connector.py#L115

A datamodule's dataloader methods have precedence over the once defined in the LightningModule, but the LightningModule code should not be altered. The user does not know that this happens, and after training is complete, the user may wishes to continue using the model instance.

Pitch

Store the dataloader references in the trainer (or data connector) directly, without "attaching" them to the user's model.
This would also enable typing inference as mentioned by @gianscarpe.

Alternatives

Keep as is, but user will not be happy.
It's also harder to debug the way it is right now.

[gianscarpe]

Hi @awaelchli, I could work on this, as I already explores datamodules for some issue I got in my own project releated to before_batch_transfer and mypy type checking :)

[awaelchli]

Aha!! Awesome, you are welcome to work on this. Ping me if you encounter any troubles along the way.
cc @PyTorchLightning/core-contributors :)

[gianscarpe]

Hi @awaelchli, I'm working on the issue and I opened a draft PR #6103 . I have some questions:

• If I pass a datamodule with only train dataloaders implemented, while I have a model with val and test dataloaders implemented, do we "combine" them? From tests it seems that the answer should be "yes", since for example predict mode is tested using ClassifierModel with train_dataloader implemented and passing a datamodule with only predict_datamodule
• If I pass a datamodule with trainer_dataloader and I have a model with the same function implemented, which one should be considered? I believe datamodule implementation

[awaelchli]

If I pass a datamodule with only train dataloaders implemented, while I have a model with val and test dataloaders implemented, do we "combine" them?

Yes it looks like combining is the current behavior. Makes sense. Even better would be to log which ones are used.

If I pass a datamodule with trainer_dataloader and I have a model with the same function implemented, which one should be considered? I believe datamodule implementation

datamodule and dataloaders passed to fit have precedence over the methods defined in the model.

[ananthsub]

@awaelchli @carmocca why do we need to bind the dataloader functions to the model at all? why can't we set these as attributes of the trainer? and then delete corresponding at the end of the call? or just call the right one of model/datamodule/dataloader inside of the training/evaluation/predict loop?

[awaelchli]

why do we need to bind the dataloader functions to the model at all?

I don't know why. It was there from the start. I don't see a reason why it should be necessary.

why can't we set these as attributes of the trainer?
we can, and it's exactly what we propose in this issue.

If @gianscarpe doesn't have time to do it, I can try to get this started in 1.4. It will be a relatively wide-spreading refactor.

[kaushikb11]

@awaelchli Adding you as an assignee as well! 🚀