Is the use of data methods on LightningModule standard?

[jhrmnn]

It took me quite some time playing with PyTorch Lightning and reading through the docs before I realized I can attach the data-related methods ([train|val]_dataloader, prepare_data) directly to LightningModule and then pass only the module to the trainer, rather than having to create a separate LightningDataModule class or passing the data loaders explicitly. When thinking of LightningModule as a task, rather than a model, this makes actually perfect sense. After all, things like batch size and learning rate are inherently related. But since I haven't seen it really explicitly documented, I wonder whether the use of data methods directly on the LightningModule is considered standard?

[awaelchli]

Hi, could you explain what you mean by

I realized I can attach the data-related methods ([train|val]_dataloader, prepare_data) directly to LightningModule

Do you mean dynamically bind the dataloader method of datamodule to lightning module?
I would not personally recommend that because it can be harder to read and debug, especially if your LightningModule also defines dataloader methods that you may expect to run in a different context.

I recommend passing the datamodule/dataloader into the fit function like so:

trainer.fit(model, datamodule=datamodule)
# or
trainer.fit(model, train_dataloader=...)

Then it will be clear to any reader that the datamodule loaders are used for training.

[jhrmnn]

Sorry I wasn't clear at all by using the word "attach". I meant statically defining the data methods as part of the LightningModule definition. So:

class MyTask(LightningModule):
    ...

    def train_dataloader(self):
        ...

    def val_dataloader(self):
        ...


task = MyTask()
trainer.fit(task)

This certainly works and the documentation even hints at it at some points. For example, in learning rate finder it says

If the model has a predefined train_dataloader method, this will be skipped.

But I haven't seen it documented in a promiment way, hence my question whether it's considered standard.

[awaelchli]

If the dataloading is tied to the task then yes it belongs there for sure. But often in practice we find that the dataloading logic can be separate from the task/lightningmodule and generalizes to other tasks. Therefore we often go for a datamodule because it's going to be reusable.

[jhrmnn]

Ok, so Lightning is rather agnostic about this, and it's simply up to me how I want to factor the code. Thanks!

[awaelchli]

Yes definitely. The only real thing missing from datamodule I think is the ability to save hyperparameters easily there. I'm pretty sure this will come with 1.4 because many have asked for this :)

[jhrmnn]

Yes, that's definitely something that I would also appreciate.