Repository: Lightning-AI/pytorch-lightning

Issue:
Improve Design of DataModule
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

### 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:

```python
class DataHooks:

    def prepare_data(...):
    """ docs... """

    def setup(...):
    """ docs... """

    def train_dataloader(...):
    """ docs... """

    etc.
```

and then in the other classes, inherit from it:

```python
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?


Discussion:
yes, just make them abstract...
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.

> 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?


Resolving patch:
diff --git a/pytorch_lightning/core/datamodule.py b/pytorch_lightning/core/datamodule.py
--- a/pytorch_lightning/core/datamodule.py
+++ b/pytorch_lightning/core/datamodule.py
@@ -21,7 +21,8 @@
 import torch
 from torch.utils.data import DataLoader
 
-from pytorch_lightning.utilities import parsing, rank_zero_only, rank_zero_warn
+from pytorch_lightning.core.hooks import DataHooks
+from pytorch_lightning.utilities import parsing, rank_zero_only
 
 
 class _DataModuleWrapper(type):
@@ -87,7 +88,7 @@ def wrapped_fn(*args, **kwargs):
     return wrapped_fn
 
 
-class LightningDataModule(object, metaclass=_DataModuleWrapper):  # pragma: no cover
+class LightningDataModule(DataHooks, metaclass=_DataModuleWrapper):
     """
     A DataModule standardizes the training, val, test splits, data preparation and transforms.
     The main advantage is consistent data splits, data preparation and transforms across models.
@@ -215,147 +216,27 @@ def has_setup_test(self):
 
     @abstractmethod
     def prepare_data(self, *args, **kwargs):
-        """
-        Use this to download and prepare data.
-        In distributed (GPU, TPU), this will only be called once.
-
-        .. warning:: Do not assign anything to the datamodule in this step since this will only be called on 1 GPU.
-
-        Pseudocode::
-
-            dm.prepare_data()
-            dm.setup()
-
-        Example::
-
-            def prepare_data(self):
-                download_imagenet()
-                clean_imagenet()
-                cache_imagenet()
-        """
+        pass
 
     @abstractmethod
     def setup(self, stage: Optional[str] = None):
-        """
-        Use this to load your data from file, split it, etc. You are safe to make state assignments here.
-        This hook is called on every process when using DDP.
-
-        Example::
-
-            def setup(self, stage):
-                data = load_data(...)
-                self.train_ds, self.val_ds, self.test_ds = split_data(data)
-        """
+        pass
 
     @abstractmethod
     def train_dataloader(self, *args, **kwargs) -> DataLoader:
-        """
-        Implement a PyTorch DataLoader for training.
-        Return:
-            Single PyTorch :class:`~torch.utils.data.DataLoader`.
-        Note:
-            Lightning adds the correct sampler for distributed and arbitrary hardware.
-            There is no need to set it yourself.
-
-        Example::
-
-            def train_dataloader(self):
-                dataset = MNIST(root=PATH, train=True, transform=transforms.ToTensor(), download=False)
-                loader = torch.utils.data.DataLoader(dataset=dataset)
-                return loader
-
-        """
-        rank_zero_warn('`train_dataloader` must be implemented to be used with the Lightning Trainer')
+        pass
 
     @abstractmethod
     def val_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
-        r"""
-        Implement a PyTorch DataLoader for training.
-        Return:
-            Single PyTorch :class:`~torch.utils.data.DataLoader`.
-        Note:
-            Lightning adds the correct sampler for distributed and arbitrary hardware.
-            There is no need to set it yourself.
-        Note:
-            You can also return a list of DataLoaders
-
-        Example::
-
-            def val_dataloader(self):
-                dataset = MNIST(root=PATH, train=False, transform=transforms.ToTensor(), download=False)
-                loader = torch.utils.data.DataLoader(dataset=dataset, shuffle=False)
-                return loader
-        """
+        pass
 
     @abstractmethod
     def test_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
-        r"""
-        Implement a PyTorch DataLoader for training.
-        Return:
-            Single PyTorch :class:`~torch.utils.data.DataLoader`.
-        Note:
-            Lightning adds the correct sampler for distributed and arbitrary hardware.
-            There is no need to set it yourself.
-        Note:
-            You can also return a list of DataLoaders
-
-        Example::
-
-            def test_dataloader(self):
-                dataset = MNIST(root=PATH, train=False, transform=transforms.ToTensor(), download=False)
-                loader = torch.utils.data.DataLoader(dataset=dataset, shuffle=False)
-                return loader
-        """
+        pass
 
     @abstractmethod
     def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
-        """
-        Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors
-        wrapped in a custom data structure.
-
-        The data types listed below (and any arbitrary nesting of them) are supported out of the box:
-
-        - :class:`torch.Tensor` or anything that implements `.to(...)`
-        - :class:`list`
-        - :class:`dict`
-        - :class:`tuple`
-        - :class:`torchtext.data.batch.Batch`
-
-        For anything else, you need to define how the data is moved to the target device (CPU, GPU, TPU, ...).
-
-        Example::
-
-            def transfer_batch_to_device(self, batch, device)
-                if isinstance(batch, CustomBatch):
-                    # move all tensors in your custom data structure to the device
-                    batch.samples = batch.samples.to(device)
-                    batch.targets = batch.targets.to(device)
-                else:
-                    batch = super().transfer_batch_to_device(data, device)
-                return batch
-
-        Args:
-            batch: A batch of data that needs to be transferred to a new device.
-            device: The target device as defined in PyTorch.
-
-        Returns:
-            A reference to the data on the new device.
-
-        Note:
-            This hook should only transfer the data and not modify it, nor should it move the data to
-            any other device than the one passed in as argument (unless you know what you are doing).
-
-        Note:
-            This hook only runs on single GPU training (no data-parallel). If you need multi-GPU support
-            for your custom batch objects, you need to define your custom
-            :class:`~torch.nn.parallel.DistributedDataParallel` or
-            :class:`~pytorch_lightning.overrides.data_parallel.LightningDistributedDataParallel` and
-            override :meth:`~pytorch_lightning.core.lightning.LightningModule.configure_ddp`.
-
-        See Also:
-            - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`
-            - :func:`~pytorch_lightning.utilities.apply_func.apply_to_collection`
-        """
+        pass
 
     @classmethod
     def add_argparse_args(cls, parent_parser: ArgumentParser) -> ArgumentParser:
diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py
--- a/pytorch_lightning/core/hooks.py
+++ b/pytorch_lightning/core/hooks.py
@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+from typing import Any, Union, List
 
 import torch
 from torch import Tensor
-from torch.nn import Module
 from torch.optim.optimizer import Optimizer
+from torch.utils.data import DataLoader
 
-from pytorch_lightning.utilities import move_data_to_device, AMPType
+from pytorch_lightning.utilities import move_data_to_device, AMPType, rank_zero_warn
 
 try:
     from apex import amp
@@ -27,7 +27,7 @@
     amp = None
 
 
-class ModelHooks(Module):
+class ModelHooks:
 
     def setup(self, stage: str):
         """
@@ -319,6 +319,214 @@ def amp_scale_loss(self, unscaled_loss, optimizer, optimizer_idx, amp_backend: A
 
         return scaled_loss
 
+
+class DataHooks:
+
+    def prepare_data(self) -> None:
+        """
+        Use this to download and prepare data.
+
+        .. warning:: DO NOT set state to the model (use `setup` instead)
+            since this is NOT called on every GPU in DDP/TPU
+
+        Example::
+
+            def prepare_data(self):
+                # good
+                download_data()
+                tokenize()
+                etc()
+
+                # bad
+                self.split = data_split
+                self.some_state = some_other_state()
+
+        In DDP prepare_data can be called in two ways (using Trainer(prepare_data_per_node)):
+
+        1. Once per node. This is the default and is only called on LOCAL_RANK=0.
+        2. Once in total. Only called on GLOBAL_RANK=0.
+
+        Example::
+
+            # DEFAULT
+            # called once per node on LOCAL_RANK=0 of that node
+            Trainer(prepare_data_per_node=True)
+
+            # call on GLOBAL_RANK=0 (great for shared file systems)
+            Trainer(prepare_data_per_node=False)
+
+        This is called before requesting the dataloaders:
+
+        .. code-block:: python
+
+            model.prepare_data()
+                if ddp/tpu: init()
+            model.setup(stage)
+            model.train_dataloader()
+            model.val_dataloader()
+            model.test_dataloader()
+        """
+
+    def train_dataloader(self) -> DataLoader:
+        """
+        Implement a PyTorch DataLoader for training.
+
+        Return:
+            Single PyTorch :class:`~torch.utils.data.DataLoader`.
+
+        The dataloader you return will not be called every epoch unless you set
+        :paramref:`~pytorch_lightning.trainer.Trainer.reload_dataloaders_every_epoch` to ``True``.
+
+        For data processing use the following pattern:
+
+            - download in :meth:`prepare_data`
+            - process and split in :meth:`setup`
+
+        However, the above are only necessary for distributed processing.
+
+        .. warning:: do not assign state in prepare_data
+
+        - :meth:`~pytorch_lightning.trainer.Trainer.fit`
+        - ...
+        - :meth:`prepare_data`
+        - :meth:`setup`
+        - :meth:`train_dataloader`
+
+        Note:
+            Lightning adds the correct sampler for distributed and arbitrary hardware.
+            There is no need to set it yourself.
+
+        Example:
+            .. code-block:: python
+
+                def train_dataloader(self):
+                    transform = transforms.Compose([transforms.ToTensor(),
+                                                    transforms.Normalize((0.5,), (1.0,))])
+                    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform,
+                                    download=True)
+                    loader = torch.utils.data.DataLoader(
+                        dataset=dataset,
+                        batch_size=self.batch_size,
+                        shuffle=True
+                    )
+                    return loader
+
+        """
+        rank_zero_warn('`train_dataloader` must be implemented to be used with the Lightning Trainer')
+
+    def test_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
+        r"""
+        Implement one or multiple PyTorch DataLoaders for testing.
+
+        The dataloader you return will not be called every epoch unless you set
+        :paramref:`~pytorch_lightning.trainer.Trainer.reload_dataloaders_every_epoch` to ``True``.
+
+        For data processing use the following pattern:
+
+            - download in :meth:`prepare_data`
+            - process and split in :meth:`setup`
+
+        However, the above are only necessary for distributed processing.
+
+        .. warning:: do not assign state in prepare_data
+
+
+        - :meth:`~pytorch_lightning.trainer.Trainer.fit`
+        - ...
+        - :meth:`prepare_data`
+        - :meth:`setup`
+        - :meth:`train_dataloader`
+        - :meth:`val_dataloader`
+        - :meth:`test_dataloader`
+
+        Note:
+            Lightning adds the correct sampler for distributed and arbitrary hardware.
+            There is no need to set it yourself.
+
+        Return:
+            Single or multiple PyTorch DataLoaders.
+
+        Example:
+            .. code-block:: python
+
+                def test_dataloader(self):
+                    transform = transforms.Compose([transforms.ToTensor(),
+                                                    transforms.Normalize((0.5,), (1.0,))])
+                    dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform,
+                                    download=True)
+                    loader = torch.utils.data.DataLoader(
+                        dataset=dataset,
+                        batch_size=self.batch_size,
+                        shuffle=False
+                    )
+
+                    return loader
+
+                # can also return multiple dataloaders
+                def test_dataloader(self):
+                    return [loader_a, loader_b, ..., loader_n]
+
+        Note:
+            If you don't need a test dataset and a :meth:`test_step`, you don't need to implement
+            this method.
+
+        Note:
+            In the case where you return multiple test dataloaders, the :meth:`test_step`
+            will have an argument ``dataloader_idx`` which matches the order here.
+        """
+
+    def val_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
+        r"""
+        Implement one or multiple PyTorch DataLoaders for validation.
+
+        The dataloader you return will not be called every epoch unless you set
+        :paramref:`~pytorch_lightning.trainer.Trainer.reload_dataloaders_every_epoch` to ``True``.
+
+        It's recommended that all data downloads and preparation happen in :meth:`prepare_data`.
+
+        - :meth:`~pytorch_lightning.trainer.Trainer.fit`
+        - ...
+        - :meth:`prepare_data`
+        - :meth:`train_dataloader`
+        - :meth:`val_dataloader`
+        - :meth:`test_dataloader`
+
+        Note:
+            Lightning adds the correct sampler for distributed and arbitrary hardware
+            There is no need to set it yourself.
+
+        Return:
+            Single or multiple PyTorch DataLoaders.
+
+        Examples:
+            .. code-block:: python
+
+                def val_dataloader(self):
+                    transform = transforms.Compose([transforms.ToTensor(),
+                                                    transforms.Normalize((0.5,), (1.0,))])
+                    dataset = MNIST(root='/path/to/mnist/', train=False,
+                                    transform=transform, download=True)
+                    loader = torch.utils.data.DataLoader(
+                        dataset=dataset,
+                        batch_size=self.batch_size,
+                        shuffle=False
+                    )
+
+                    return loader
+
+                # can also return multiple dataloaders
+                def val_dataloader(self):
+                    return [loader_a, loader_b, ..., loader_n]
+
+        Note:
+            If you don't need a validation dataset and a :meth:`validation_step`, you don't need to
+            implement this method.
+
+        Note:
+            In the case where you return multiple validation dataloaders, the :meth:`validation_step`
+            will have an argument ``dataloader_idx`` which matches the order here.
+        """
+
     def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
         """
         Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -17,7 +17,7 @@
 import os
 import re
 import tempfile
-from abc import ABC, abstractmethod
+from abc import ABC
 from argparse import Namespace
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
@@ -27,11 +27,10 @@
 from torch.nn import Module
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim.optimizer import Optimizer
-from torch.utils.data import DataLoader
 
 from pytorch_lightning import _logger as log
 from pytorch_lightning.core.grads import GradInformation
-from pytorch_lightning.core.hooks import ModelHooks
+from pytorch_lightning.core.hooks import ModelHooks, DataHooks
 from pytorch_lightning.core.memory import ModelSummary
 from pytorch_lightning.core.saving import ALLOWED_CONFIG_TYPES, PRIMITIVE_TYPES, ModelIO
 from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel
@@ -48,7 +47,7 @@
     XLA_AVAILABLE = True
 
 
-class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, ModelHooks, Module):
+class LightningModule(ABC, DeviceDtypeModuleMixin, GradInformation, ModelIO, ModelHooks, DataHooks, Module):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -1221,224 +1220,6 @@ def tbptt_split_batch(self, batch, split_size):
 
         return splits
 
-    def prepare_data(self) -> None:
-        """
-        Use this to download and prepare data.
-
-        .. warning:: DO NOT set state to the model (use `setup` instead)
-            since this is NOT called on every GPU in DDP/TPU
-
-        Example::
-
-            def prepare_data(self):
-                # good
-                download_data()
-                tokenize()
-                etc()
-
-                # bad
-                self.split = data_split
-                self.some_state = some_other_state()
-
-        In DDP prepare_data can be called in two ways (using Trainer(prepare_data_per_node)):
-
-        1. Once per node. This is the default and is only called on LOCAL_RANK=0.
-        2. Once in total. Only called on GLOBAL_RANK=0.
-
-        Example::
-
-            # DEFAULT
-            # called once per node on LOCAL_RANK=0 of that node
-            Trainer(prepare_data_per_node=True)
-
-            # call on GLOBAL_RANK=0 (great for shared file systems)
-            Trainer(prepare_data_per_node=False)
-
-        This is called before requesting the dataloaders:
-
-        .. code-block:: python
-
-            model.prepare_data()
-                if ddp/tpu: init()
-            model.setup(stage)
-            model.train_dataloader()
-            model.val_dataloader()
-            model.test_dataloader()
-        """
-
-    def train_dataloader(self) -> DataLoader:
-        """
-        Implement a PyTorch DataLoader for training.
-
-        Return:
-            Single PyTorch :class:`~torch.utils.data.DataLoader`.
-
-        The dataloader you return will not be called every epoch unless you set
-        :paramref:`~pytorch_lightning.trainer.Trainer.reload_dataloaders_every_epoch` to ``True``.
-
-        For data processing use the following pattern:
-
-            - download in :meth:`prepare_data`
-            - process and split in :meth:`setup`
-
-        However, the above are only necessary for distributed processing.
-
-        .. warning:: do not assign state in prepare_data
-
-        - :meth:`~pytorch_lightning.trainer.Trainer.fit`
-        - ...
-        - :meth:`prepare_data`
-        - :meth:`setup`
-        - :meth:`train_dataloader`
-
-        Note:
-            Lightning adds the correct sampler for distributed and arbitrary hardware.
-            There is no need to set it yourself.
-
-        Example:
-            .. code-block:: python
-
-                def train_dataloader(self):
-                    transform = transforms.Compose([transforms.ToTensor(),
-                                                    transforms.Normalize((0.5,), (1.0,))])
-                    dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform,
-                                    download=True)
-                    loader = torch.utils.data.DataLoader(
-                        dataset=dataset,
-                        batch_size=self.batch_size,
-                        shuffle=True
-                    )
-                    return loader
-
-        """
-        rank_zero_warn('`train_dataloader` must be implemented to be used with the Lightning Trainer')
-
-    def tng_dataloader(self):  # todo: remove in v1.0.0
-        """
-        Warnings:
-            Deprecated in v0.5.0. Use :meth:`train_dataloader` instead. Will be removed in 1.0.0.
-        """
-        output = self.train_dataloader()
-        rank_zero_warn(
-            "`tng_dataloader` has been renamed to `train_dataloader` since v0.5.0."
-            " and this method will be removed in v1.0.0",
-            DeprecationWarning,
-        )
-        return output
-
-    def test_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
-        r"""
-        Implement one or multiple PyTorch DataLoaders for testing.
-
-        The dataloader you return will not be called every epoch unless you set
-        :paramref:`~pytorch_lightning.trainer.Trainer.reload_dataloaders_every_epoch` to ``True``.
-
-        For data processing use the following pattern:
-
-            - download in :meth:`prepare_data`
-            - process and split in :meth:`setup`
-
-        However, the above are only necessary for distributed processing.
-
-        .. warning:: do not assign state in prepare_data
-
-
-        - :meth:`~pytorch_lightning.trainer.Trainer.fit`
-        - ...
-        - :meth:`prepare_data`
-        - :meth:`setup`
-        - :meth:`train_dataloader`
-        - :meth:`val_dataloader`
-        - :meth:`test_dataloader`
-
-        Note:
-            Lightning adds the correct sampler for distributed and arbitrary hardware.
-            There is no need to set it yourself.
-
-        Return:
-            Single or multiple PyTorch DataLoaders.
-
-        Example:
-            .. code-block:: python
-
-                def test_dataloader(self):
-                    transform = transforms.Compose([transforms.ToTensor(),
-                                                    transforms.Normalize((0.5,), (1.0,))])
-                    dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform,
-                                    download=True)
-                    loader = torch.utils.data.DataLoader(
-                        dataset=dataset,
-                        batch_size=self.batch_size,
-                        shuffle=False
-                    )
-
-                    return loader
-
-                # can also return multiple dataloaders
-                def test_dataloader(self):
-                    return [loader_a, loader_b, ..., loader_n]
-
-        Note:
-            If you don't need a test dataset and a :meth:`test_step`, you don't need to implement
-            this method.
-
-        Note:
-            In the case where you return multiple test dataloaders, the :meth:`test_step`
-            will have an argument ``dataloader_idx`` which matches the order here.
-        """
-
-    def val_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
-        r"""
-        Implement one or multiple PyTorch DataLoaders for validation.
-
-        The dataloader you return will not be called every epoch unless you set
-        :paramref:`~pytorch_lightning.trainer.Trainer.reload_dataloaders_every_epoch` to ``True``.
-
-        It's recommended that all data downloads and preparation happen in :meth:`prepare_data`.
-
-        - :meth:`~pytorch_lightning.trainer.Trainer.fit`
-        - ...
-        - :meth:`prepare_data`
-        - :meth:`train_dataloader`
-        - :meth:`val_dataloader`
-        - :meth:`test_dataloader`
-
-        Note:
-            Lightning adds the correct sampler for distributed and arbitrary hardware
-            There is no need to set it yourself.
-
-        Return:
-            Single or multiple PyTorch DataLoaders.
-
-        Examples:
-            .. code-block:: python
-
-                def val_dataloader(self):
-                    transform = transforms.Compose([transforms.ToTensor(),
-                                                    transforms.Normalize((0.5,), (1.0,))])
-                    dataset = MNIST(root='/path/to/mnist/', train=False,
-                                    transform=transform, download=True)
-                    loader = torch.utils.data.DataLoader(
-                        dataset=dataset,
-                        batch_size=self.batch_size,
-                        shuffle=False
-                    )
-
-                    return loader
-
-                # can also return multiple dataloaders
-                def val_dataloader(self):
-                    return [loader_a, loader_b, ..., loader_n]
-
-        Note:
-            If you don't need a validation dataset and a :meth:`validation_step`, you don't need to
-            implement this method.
-
-        Note:
-            In the case where you return multiple validation dataloaders, the :meth:`validation_step`
-            will have an argument ``dataloader_idx`` which matches the order here.
-        """
-
     def summarize(self, mode: str = ModelSummary.MODE_DEFAULT) -> ModelSummary:
         model_summary = ModelSummary(self, mode=mode)
         log.info('\n' + str(model_summary))


Test patch: