remove dataloader patching on the LightningModule

CHANGELOG.md

@@ -466,6 +466,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Removed `should_rank_save_checkpoint` property from Trainer ([#9433](https://github.com/PyTorchLightning/pytorch-lightning/pull/9433))
 
 
+- Removed automatic patching of `{train,val,test,predict}_dataloader()` on the `LightningModule` ([#9764](https://github.com/PyTorchLightning/pytorch-lightning/pull/9764))
+
+
 ### Fixed
 
 
@@ -525,6 +528,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Reset `val_dataloader` in `tuner/batch_size_scaling` ([#9857](https://github.com/PyTorchLightning/pytorch-lightning/pull/9857))
 
 
+- Fixed undesired side effects being caused by `Trainer` patching dataloader methods on the `LightningModule` ([#9764](https://github.com/PyTorchLightning/pytorch-lightning/pull/9764))
+
+
 ## [1.4.9] - 2021-09-30
 
 - Fixed `lr_find` to generate same results on multiple calls ([#9704](https://github.com/PyTorchLightning/pytorch-lightning/pull/9704))

pytorch_lightning/plugins/plugins_registry.py

@@ -127,11 +127,7 @@ def is_register_plugins_overridden(plugin: type) -> bool:
     else:
         return False
 
-    if hasattr(plugin_attr, "patch_loader_code"):
-        is_overridden = plugin_attr.patch_loader_code != str(super_attr.__code__)
-    else:
-        is_overridden = plugin_attr.__code__ is not super_attr.__code__
-    return is_overridden
+    return plugin_attr.__code__ is not super_attr.__code__
 
 
 def call_training_type_register_plugins(root: Path, base_module: str) -> None:

pytorch_lightning/plugins/training_type/deepspeed.py

@@ -589,8 +589,9 @@ def _auto_select_batch_size(self):
         # train_micro_batch_size_per_gpu is used for throughput logging purposes
         # by default we try to use the batch size of the loader
         batch_size = 1
-        if hasattr(self.lightning_module, "train_dataloader"):
-            train_dataloader = self.lightning_module.train_dataloader()
+        train_dl_source = self.lightning_module.trainer.data_connector._train_dataloader_source
+        if train_dl_source.is_available():
+            train_dataloader = train_dl_source.request()
             if hasattr(train_dataloader, "batch_sampler"):
                 batch_size = train_dataloader.batch_sampler.batch_size
         return batch_size

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -19,14 +19,13 @@
 
 import torch
 import torch.multiprocessing as mp
-from torch.nn import Module
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.training_type.ddp_spawn import DDPSpawnPlugin
-from pytorch_lightning.trainer.connectors.data_connector import _PatchDataLoader
+from pytorch_lightning.trainer.connectors.data_connector import DataConnector
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities import (
     _OMEGACONF_AVAILABLE,
@@ -95,19 +94,18 @@ def _validate_dataloader(dataloaders: Union[List[DataLoader], DataLoader]) -> No
                 )
 
     @staticmethod
-    def _validate_patched_dataloaders(model: Module) -> None:
+    def _validate_patched_dataloaders(model: "pl.LightningModule") -> None:
         """Validate and fail fast if the dataloaders were passed directly to fit."""
-        if hasattr(model, "train_dataloader") and isinstance(model.train_dataloader, _PatchDataLoader):
-            TPUSpawnPlugin._validate_dataloader(model.train_dataloader.dataloader)
-
-        if hasattr(model, "val_dataloader") and isinstance(model.val_dataloader, _PatchDataLoader):
-            TPUSpawnPlugin._validate_dataloader(model.val_dataloader.dataloader)
-
-        if hasattr(model, "test_dataloader") and isinstance(model.test_dataloader, _PatchDataLoader):
-            TPUSpawnPlugin._validate_dataloader(model.test_dataloader.dataloader)
-
-        if hasattr(model, "predict_dataloader") and isinstance(model.predict_dataloader, _PatchDataLoader):
-            TPUSpawnPlugin._validate_dataloader(model.predict_dataloader.dataloader)
+        connector: DataConnector = model.trainer.data_connector
+        sources = (
+            connector._train_dataloader_source,
+            connector._val_dataloader_source,
+            connector._test_dataloader_source,
+            connector._predict_dataloader_source,
+        )
+        for dataloader_source in sources:
+            if not dataloader_source.is_module():
+                TPUSpawnPlugin._validate_dataloader(dataloader_source.instance)
 
     def connect(self, model: "pl.LightningModule") -> None:
         TPUSpawnPlugin._validate_patched_dataloaders(model)

pytorch_lightning/trainer/configuration_validator.py

@@ -67,7 +67,7 @@ def __verify_train_loop_configuration(self, model: "pl.LightningModule") -> None
         # -----------------------------------
         # verify model has a train dataloader
         # -----------------------------------
-        has_train_dataloader = is_overridden("train_dataloader", model)
+        has_train_dataloader = self.trainer.data_connector._train_dataloader_source.is_available()
         if not has_train_dataloader:
             raise MisconfigurationException(
                 "No `train_dataloader()` method defined. Lightning `Trainer` expects as minimum a"
@@ -176,7 +176,7 @@ def __verify_eval_loop_configuration(self, model: "pl.LightningModule", stage: s
             )
 
     def __verify_predict_loop_configuration(self, model: "pl.LightningModule") -> None:
-        has_predict_dataloader = is_overridden("predict_dataloader", model)
+        has_predict_dataloader = self.trainer.data_connector._predict_dataloader_source.is_available()
         if not has_predict_dataloader:
             raise MisconfigurationException("Dataloader not found for `Trainer.predict`")
         # ----------------------------------------------

pytorch_lightning/trainer/connectors/data_connector.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from dataclasses import dataclass
 from functools import partial
-from typing import Callable, Iterable, Optional, Union
+from typing import Iterable, Optional, Union
 
 import pytorch_lightning as pl
 from pytorch_lightning.utilities import rank_zero_deprecation
@@ -47,6 +48,11 @@ def __init__(
         self.test_data_fetcher = test_data_fetcher
         self.sanity_check_data_fetcher: Optional[AbstractDataFetcher] = None
 
+        self._train_dataloader_source = _DataLoaderSource()
+        self._val_dataloader_source = _DataLoaderSource()
+        self._test_dataloader_source = _DataLoaderSource()
+        self._predict_dataloader_source = _DataLoaderSource()
+
     @property
     def evaluation_data_fetcher(self) -> Optional[AbstractDataFetcher]:
         if self.trainer.sanity_checking:
@@ -190,27 +196,23 @@ def attach_dataloaders(
         test_dataloaders: Optional[EVAL_DATALOADERS] = None,
         predict_dataloaders: Optional[EVAL_DATALOADERS] = None,
     ) -> None:
-        # when dataloader is passed via fit, patch the train_dataloader
-        # functions to overwrite with these implementations
-        if train_dataloaders is not None:
-            self.trainer.train_dataloader = None
-            train_dataloader = _PatchDataLoader(train_dataloaders, "train")
-            train_dataloader.patch(model)
-
-        if val_dataloaders is not None:
-            self.trainer.val_dataloaders = None
-            val_dataloader = _PatchDataLoader(val_dataloaders, "val")
-            val_dataloader.patch(model)
-
-        if test_dataloaders is not None:
-            self.trainer.test_dataloaders = None
-            test_dataloader = _PatchDataLoader(test_dataloaders, "test")
-            test_dataloader.patch(model)
-
-        if predict_dataloaders is not None:
-            self.trainer.predict_dataloaders = None
-            predict_dataloader = _PatchDataLoader(predict_dataloaders, "predict")
-            predict_dataloader.patch(model)
+        self.trainer.train_dataloader = None
+        self.trainer.val_dataloaders = None
+        self.trainer.test_dataloaders = None
+        self.trainer.predict_dataloaders = None
+
+        self._train_dataloader_source = _DataLoaderSource(
+            train_dataloaders if train_dataloaders is not None else model, "train_dataloader"
+        )
+        self._val_dataloader_source = _DataLoaderSource(
+            val_dataloaders if val_dataloaders is not None else model, "val_dataloader"
+        )
+        self._test_dataloader_source = _DataLoaderSource(
+            test_dataloaders if test_dataloaders is not None else model, "test_dataloader"
+        )
+        self._predict_dataloader_source = _DataLoaderSource(
+            predict_dataloaders if predict_dataloaders is not None else model, "predict_dataloader"
+        )
 
     def attach_datamodule(
         self, model: "pl.LightningModule", datamodule: Optional["pl.LightningDataModule"] = None
@@ -219,11 +221,10 @@ def attach_datamodule(
         if datamodule is None:
             return
 
-        # Override loader hooks
-        dl_methods = ("train_dataloader", "val_dataloader", "test_dataloader", "predict_dataloader")
-        for method in dl_methods:
-            if is_overridden(method, datamodule):
-                setattr(model, method, getattr(datamodule, method))
+        self._train_dataloader_source = _DataLoaderSource(datamodule, "train_dataloader")
+        self._val_dataloader_source = _DataLoaderSource(datamodule, "val_dataloader")
+        self._test_dataloader_source = _DataLoaderSource(datamodule, "test_dataloader")
+        self._predict_dataloader_source = _DataLoaderSource(datamodule, "predict_dataloader")
 
         # Override data transfer hooks if dataset-specific to_device logic has been defined in datamodule
         batch_transfer_hooks = ("on_before_batch_transfer", "transfer_batch_to_device", "on_after_batch_transfer")
@@ -238,13 +239,6 @@ def attach_datamodule(
         if hasattr(datamodule, "data_pipeline"):
             model.data_pipeline = datamodule.data_pipeline
 
-    @staticmethod
-    def detach_data(model: "pl.LightningModule") -> None:
-        for stage in ("train", "val", "test", "predict"):
-            loader = getattr(model, f"{stage}_dataloader", None)
-            if isinstance(loader, _PatchDataLoader):
-                loader.unpatch(model)
-
     def teardown(self) -> None:
         if self.train_data_fetcher:
             self.train_data_fetcher.teardown()
@@ -260,32 +254,58 @@ def teardown(self) -> None:
             self.sanity_check_data_fetcher = None
 
 
-class _PatchDataLoader:
-    r"""
-    Callable object for patching dataloaders passed into trainer.fit().
-    Use this class to override model.*_dataloader() and be pickle-compatible.
+@dataclass
+class _DataLoaderSource:
+    """Stores the information where the dataloaders come from.
+
+    The source can be
 
-    Args:
-        dataloader: Dataloader object to return when called.
+    1. from a ``*_datalaoder()`` method on the :class:`~pytorch_lightning.core.lightning.LightningModule`,
+    2. from a ``*_datalaoder()`` method on the :class:`~pytorch_lightning.core.datamodule.LightningDataModule`,
+    3. a direct instance of a :class:`~torch.utils.data.DataLoader` or supported collections thereof.
+
+    Arguments:
+        instance: A LightningModule, LightningDataModule, or (a collection of) dataloader(s).
+        name: A name for this dataloader source. If the instance is a module, the name corresponds to the hook
+            that returns the desired dataloader(s).
     """
 
-    def __init__(self, dataloader: Union[TRAIN_DATALOADERS, EVAL_DATALOADERS], stage: str) -> None:
-        self.dataloader = dataloader
+    instance: Optional[
+        Union[TRAIN_DATALOADERS, EVAL_DATALOADERS, "pl.LightningModule", "pl.LightningDataModule"]
+    ] = None
+    name: str = ""
+
+    def request(self) -> Union[TRAIN_DATALOADERS, EVAL_DATALOADERS]:
+        """Returns the dataloader from the source.
+
+        If the source is a module, the method with the corresponding
+        :attr:`name` gets called.
+        """
+        from pytorch_lightning import LightningDataModule, LightningModule  # prevent cyclic import
+
+        if not self.name:
+            return self.instance
+
+        if isinstance(self.instance, LightningModule):
+            return self.instance.trainer.call_hook(self.name, pl_module=self.instance)
+
+        if isinstance(self.instance, LightningDataModule):
+            return getattr(self.instance, self.name)()
+
+        return self.instance
+
+    def is_available(self) -> bool:
+        """Returns whether the source dataloader is available.
 
-        # cannot pickle __code__ so cannot verify if PatchDataloader
-        # exists which shows dataloader methods have been overwritten.
-        # so, we hack it by using the string representation
-        self.patch_loader_code = str(self.__call__.__code__)
-        self._old_loader: Optional[Callable] = None
-        self.stage = stage
+        If the source is a module it checks that the method with given :attr:`name` is overridden.
+        """
+        return not self.is_module() or is_overridden(self.name, self.instance)
 
-    def __call__(self) -> Union[TRAIN_DATALOADERS, EVAL_DATALOADERS]:
-        return self.dataloader
+    def is_module(self) -> bool:
+        """Returns whether the the DataLoader source is a LightningModule or a LightningDataModule.
 
-    def patch(self, model: "pl.LightningModule") -> None:
-        self._old_loader = getattr(model, self.stage + "_dataloader")
-        setattr(model, self.stage + "_dataloader", self)
+        It does not check whether ``*_dataloader`` methods are actually overridden.
+        """
+        from pytorch_lightning import LightningDataModule, LightningModule  # prevent cyclic import
 
-    def unpatch(self, model: "pl.LightningModule") -> None:
-        setattr(model, self.stage + "_dataloader", self._old_loader)
-        self._old_loader = None
+        return isinstance(self.instance, (LightningModule, LightningDataModule))

pytorch_lightning/trainer/data_loading.py

@@ -487,10 +487,10 @@ def reset_val_dataloader(self, model: Optional["pl.LightningModule"] = None) ->
         Args:
             model: The `LightningModule` if called outside of the trainer scope.
         """
+        source = self.data_connector._val_dataloader_source
         pl_module = self.lightning_module or model
-        has_loader = is_overridden("val_dataloader", pl_module)
         has_step = is_overridden("validation_step", pl_module)
-        if has_loader and has_step:
+        if source.is_available() and has_step:
             self.num_val_batches, self.val_dataloaders = self._reset_eval_dataloader(
                 RunningStage.VALIDATING, model=pl_module
             )
@@ -501,10 +501,10 @@ def reset_test_dataloader(self, model: Optional["pl.LightningModule"] = None) ->
         Args:
             model: The `LightningModule` if called outside of the trainer scope.
         """
+        source = self.data_connector._test_dataloader_source
         pl_module = self.lightning_module or model
-        has_loader = is_overridden("test_dataloader", pl_module)
         has_step = is_overridden("test_step", pl_module)
-        if has_loader and has_step:
+        if source.is_available() and has_step:
             self.num_test_batches, self.test_dataloaders = self._reset_eval_dataloader(
                 RunningStage.TESTING, model=pl_module
             )
@@ -515,9 +515,9 @@ def reset_predict_dataloader(self, model: Optional["pl.LightningModule"] = None)
         Args:
             model: The `LightningModule` if called outside of the trainer scope.
         """
+        source = self.data_connector._predict_dataloader_source
         pl_module = self.lightning_module or model
-        has_loader = is_overridden("predict_dataloader", pl_module)
-        if has_loader:
+        if source.is_available():
             self.num_predict_batches, self.predict_dataloaders = self._reset_eval_dataloader(
                 RunningStage.PREDICTING, model=pl_module
             )
@@ -539,14 +539,16 @@ def reset_train_val_dataloaders(self, model: Optional["pl.LightningModule"] = No
     def request_dataloader(
         self, stage: RunningStage, model: Optional["pl.LightningModule"] = None
     ) -> Union[DataLoader, List[DataLoader]]:
-        """Handles downloading data in the GPU or TPU case.
+        """Requests a dataloader from the given model by calling dataloader hooks corresponding to the given stage.
 
         Returns:
-            The dataloader
+            The requested dataloader
         """
+        source = getattr(self.data_connector, f"_{stage.dataloader_prefix}_dataloader_source")
+
         hook = f"{stage.dataloader_prefix}_dataloader"
         self.call_hook("on_" + hook, pl_module=model)
-        dataloader = self.call_hook(hook, pl_module=model)
+        dataloader = source.request()
         if isinstance(dataloader, tuple):
             dataloader = list(dataloader)
         self.training_type_plugin.barrier("get_dataloaders")

pytorch_lightning/trainer/trainer.py

@@ -1259,7 +1259,9 @@ def _run_predict(self) -> Optional[_PREDICT_OUTPUT]:
             return self.predict_loop.run()
 
     def _run_sanity_check(self, ref_model):
-        using_val_step = ref_model.val_dataloader is not None and is_overridden("validation_step", ref_model)
+        using_val_step = self.data_connector._val_dataloader_source.is_available() and is_overridden(
+            "validation_step", ref_model
+        )
         should_sanity_check = using_val_step and self.num_sanity_val_steps > 0 and self.limit_val_batches > 0
 
         # run tiny validation (if validation defined)
@@ -1358,8 +1360,6 @@ def _call_teardown_hook(self) -> None:
         if self.datamodule is not None:
             self.datamodule.teardown(stage=fn)
 
-        self.data_connector.detach_data(self.lightning_module)
-
         self.call_hook("teardown", stage=fn)
 
         self.lightning_module._current_fx_name = None

pytorch_lightning/tuner/batch_size_scaling.py

@@ -51,7 +51,7 @@ def scale_batch_size(
             " If this is not the intended behavior, please remove either one."
         )
 
-    if hasattr(model.train_dataloader, "patch_loader_code"):
+    if not trainer.data_connector._train_dataloader_source.is_module():
         raise MisconfigurationException(
             "The batch scaling feature cannot be used with dataloaders passed directly to `.fit()`."
             " Please disable the feature or incorporate the dataloader into the model."

pytorch_lightning/utilities/model_helpers.py

@@ -64,10 +64,7 @@ def is_overridden(
     if parent_attr is None:
         raise ValueError("The parent should define the method")
 
-    # cannot pickle `__code__` so cannot verify if `PatchDataloader`
-    # exists which shows dataloader methods have been overwritten.
-    # so, we hack it by using the string representation
-    instance_code = getattr(instance_attr, "patch_loader_code", None) or instance_attr.__code__
+    instance_code = instance_attr.__code__
     parent_code = parent_attr.__code__
 
     return instance_code != parent_code