Avoid patching common DataHooks to the LightningModule (#10603)

Author: rohitgr7

CHANGELOG.md

@@ -514,6 +514,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Removed `DeepSpeedPlugin.{precision,amp_type,amp_level}` properties ([#10657](https://github.com/PyTorchLightning/pytorch-lightning/pull/10657))
 
 
+- Removed patching of `on_before_batch_transfer`, `transfer_batch_to_device` and `on_after_batch_transfer` hooks in `LightningModule` ([#10603](https://github.com/PyTorchLightning/pytorch-lightning/pull/10603))
+
+
 - Removed argument `return_result` from the `DDPSpawnPlugin.spawn()` method ([#10867](https://github.com/PyTorchLightning/pytorch-lightning/pull/10867))
 
 

pytorch_lightning/core/lightning.py

@@ -38,6 +38,7 @@
 from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.core.saving import ModelIO
 from pytorch_lightning.loggers import LightningLoggerBase
+from pytorch_lightning.trainer.connectors.data_connector import _DataHookSelector
 from pytorch_lightning.trainer.connectors.logger_connector.fx_validator import _FxValidator
 from pytorch_lightning.utilities import _IS_WINDOWS, _TORCH_GREATER_EQUAL_1_10, GradClipAlgorithmType
 from pytorch_lightning.utilities.apply_func import apply_to_collection, convert_to_tensors
@@ -259,9 +260,16 @@ def _apply_batch_transfer_handler(
         self, batch: Any, device: Optional[torch.device] = None, dataloader_idx: int = 0
     ) -> Any:
         device = device or self.device
-        batch = self.on_before_batch_transfer(batch, dataloader_idx)
-        batch = self.transfer_batch_to_device(batch, device, dataloader_idx)
-        batch = self.on_after_batch_transfer(batch, dataloader_idx)
+        datahook_selector = (
+            _DataHookSelector(self, None) if self.trainer is None else self.trainer._data_connector._datahook_selector
+        )
+
+        hook = datahook_selector.get_hook("on_before_batch_transfer")
+        batch = hook(batch, dataloader_idx)
+        hook = datahook_selector.get_hook("transfer_batch_to_device")
+        batch = hook(batch, device, dataloader_idx)
+        hook = datahook_selector.get_hook("on_after_batch_transfer")
+        batch = hook(batch, dataloader_idx)
         return batch
 
     def print(self, *args, **kwargs) -> None:

pytorch_lightning/trainer/configuration_validator.py

@@ -206,8 +206,11 @@ def __verify_dp_batch_transfer_support(trainer: "pl.Trainer", model: "pl.Lightni
     """Raise Misconfiguration exception since these hooks are not supported in DP mode."""
     # TODO: Remove this blocker once batch transfer to device is integrated in Lightning for DP mode.
     batch_transfer_hooks = ("on_before_batch_transfer", "transfer_batch_to_device", "on_after_batch_transfer")
+    datahook_selector = trainer._data_connector._datahook_selector
     for hook in batch_transfer_hooks:
-        if trainer._accelerator_connector.use_dp and is_overridden(hook, model):
+        if trainer._accelerator_connector.use_dp and (
+            is_overridden(hook, datahook_selector.model) or is_overridden(hook, datahook_selector.datamodule)
+        ):
             raise MisconfigurationException(f"Overriding `{hook}` is not supported in DP mode.")
 
 

pytorch_lightning/trainer/connectors/data_connector.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 import multiprocessing
 import os
-from dataclasses import dataclass
-from typing import Any, Collection, List, Optional, Tuple, Union
+from dataclasses import dataclass, field
+from typing import Any, Callable, Collection, List, Optional, Tuple, Union
 from weakref import proxy
 
 from torch.utils.data import DataLoader, RandomSampler, Sampler, SequentialSampler
@@ -40,7 +40,9 @@
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation, rank_zero_warn
 from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
-from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from pytorch_lightning.utilities.warnings import PossibleUserWarning, WarningCache
+
+warning_cache = WarningCache()
 
 
 class DataConnector:
@@ -52,6 +54,8 @@ def __init__(self, trainer: "pl.Trainer", multiple_trainloader_mode: str = "max_
         self._test_dataloader_source = _DataLoaderSource(None, "")
         self._predict_dataloader_source = _DataLoaderSource(None, "")
 
+        self._datahook_selector = _DataHookSelector(None, None)
+
     @property
     def _should_reload_train_dl(self) -> bool:
         """Check if train dataloader should be reloaded."""
@@ -192,6 +196,8 @@ def attach_datamodule(
         self, model: "pl.LightningModule", datamodule: Optional["pl.LightningDataModule"] = None
     ) -> None:
         # If we have a datamodule, attach necessary hooks + dataloaders
+        self._datahook_selector = _DataHookSelector(model, datamodule)
+
         if datamodule is None:
             return
 
@@ -200,12 +206,6 @@ def attach_datamodule(
         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")
-        for hook in batch_transfer_hooks:
-            if is_overridden(hook, datamodule):
-                setattr(model, hook, getattr(datamodule, hook))
-
         self.trainer.datamodule = datamodule
         datamodule.trainer = self.trainer
 
@@ -555,3 +555,48 @@ def is_module(self) -> bool:
         from pytorch_lightning import LightningDataModule, LightningModule  # prevent cyclic import
 
         return isinstance(self.instance, (LightningModule, LightningDataModule))
+
+
+@dataclass
+class _DataHookSelector:
+    """Stores the info about the shared DataHooks within LightningModule and LightningDataModule.
+
+    The hook source can be
+
+    1. a method from the :class:`~pytorch_lightning.core.lightning.LightningModule`,
+    2. a method from the :class:`~pytorch_lightning.core.datamodule.LightningDataModule`,
+
+    Arguments:
+        model: A LightningModule
+        datamodule: A LightningDataModule
+    """
+
+    model: "pl.LightningModule"
+    datamodule: Optional["pl.LightningDataModule"]
+    _valid_hooks: Tuple[str] = field(
+        default=("on_before_batch_transfer", "transfer_batch_to_device", "on_after_batch_transfer")
+    )
+
+    def get_hook(self, hook_name: str) -> Callable:
+        if hook_name not in self._valid_hooks:
+            raise ValueError(
+                f"`{hook_name}` is not a shared hook within `LightningModule` and `LightningDataModule`."
+                f" Valid hooks are {self._valid_hooks}."
+            )
+
+        if self.datamodule is None:
+            return getattr(self.model, hook_name)
+
+        if is_overridden(hook_name, self.datamodule):
+            if is_overridden(hook_name, self.model):
+                warning_cache.warn(
+                    f"You have overridden `{hook_name}` in both `LightningModule` and `LightningDataModule`."
+                    " It will use the implementation from `LightningDataModule` instance."
+                )
+            return getattr(self.datamodule, hook_name)
+
+        warning_cache.warn(
+            f"You have overridden `{hook_name}` in `LightningModule` but have passed in a"
+            " `LightningDataModule`. It will use the implementation from `LightningModule` instance."
+        )
+        return getattr(self.model, hook_name)

tests/core/test_datamodules.py

@@ -26,7 +26,6 @@
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities import _OMEGACONF_AVAILABLE, AttributeDict
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.model_helpers import is_overridden
 from tests.helpers import BoringDataModule, BoringModel
 from tests.helpers.datamodules import ClassifDataModule
 from tests.helpers.runif import RunIf
@@ -309,15 +308,11 @@ def transfer_batch_to_device(self, batch, device, dataloader_idx):
     batch = CustomBatch((torch.zeros(5, 32), torch.ones(5, 1, dtype=torch.long)))
 
     trainer = Trainer(accelerator="gpu", devices=1)
+    model.trainer = trainer
     # running .fit() would require us to implement custom data loaders, we mock the model reference instead
     get_module_mock.return_value = model
-    if is_overridden("transfer_batch_to_device", dm):
-        model.transfer_batch_to_device = dm.transfer_batch_to_device
-
-    model.on_before_batch_transfer = dm.on_before_batch_transfer
-    model.transfer_batch_to_device = dm.transfer_batch_to_device
-    model.on_after_batch_transfer = dm.on_after_batch_transfer
 
+    trainer._data_connector.attach_datamodule(model, datamodule=dm)
     batch_gpu = trainer.strategy.batch_to_device(batch, expected_device)
 
     assert dm.on_before_batch_transfer_hook_rank == 0

tests/trainer/connectors/test_data_connector.py

@@ -17,9 +17,10 @@
 from torch.utils.data import DataLoader
 
 from pytorch_lightning import Trainer
-from pytorch_lightning.trainer.connectors.data_connector import _DataLoaderSource
+from pytorch_lightning.trainer.connectors.data_connector import _DataHookSelector, _DataLoaderSource, warning_cache
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.warnings import PossibleUserWarning
+from tests.deprecated_api import no_warning_call
 from tests.helpers import BoringDataModule, BoringModel
 from tests.helpers.boring_model import RandomDataset
 
@@ -71,6 +72,77 @@ def test_dataloader_source_request_from_module():
     module.foo.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    "hook_name", ("on_before_batch_transfer", "transfer_batch_to_device", "on_after_batch_transfer")
+)
+class TestDataHookSelector:
+    def overridden_func(self, batch, *args, **kwargs):
+        return batch
+
+    def reset_instances(self):
+        return BoringDataModule(), BoringModel(), Trainer()
+
+    def test_no_datamodule_no_overridden(self, hook_name):
+        model, _, trainer = self.reset_instances()
+        trainer._data_connector.attach_datamodule(model, datamodule=None)
+        with no_warning_call(match="have overridden `{hook_name}` in both"):
+            hook = trainer._data_connector._datahook_selector.get_hook(hook_name)
+
+        assert hook == getattr(model, hook_name)
+
+    def test_with_datamodule_no_overridden(self, hook_name):
+        model, dm, trainer = self.reset_instances()
+        trainer._data_connector.attach_datamodule(model, datamodule=dm)
+        with no_warning_call(match="have overridden `{hook_name}` in both"):
+            hook = trainer._data_connector._datahook_selector.get_hook(hook_name)
+
+        assert hook == getattr(model, hook_name)
+
+    def test_override_model_hook(self, hook_name):
+        model, dm, trainer = self.reset_instances()
+        trainer._data_connector.attach_datamodule(model, datamodule=dm)
+        with no_warning_call(match="have overridden `{hook_name}` in both"):
+            hook = trainer._data_connector._datahook_selector.get_hook(hook_name)
+
+        assert hook == getattr(model, hook_name)
+
+    def test_override_datamodule_hook(self, hook_name):
+        model, dm, trainer = self.reset_instances()
+        trainer._data_connector.attach_datamodule(model, datamodule=dm)
+        setattr(dm, hook_name, self.overridden_func)
+        with no_warning_call(match="have overridden `{hook_name}` in both"):
+            hook = trainer._data_connector._datahook_selector.get_hook(hook_name)
+
+        assert hook == getattr(dm, hook_name)
+
+    def test_override_both_model_and_datamodule(self, hook_name):
+        model, dm, trainer = self.reset_instances()
+        trainer._data_connector.attach_datamodule(model, datamodule=dm)
+        setattr(model, hook_name, self.overridden_func)
+        setattr(dm, hook_name, self.overridden_func)
+        with pytest.warns(UserWarning, match=f"have overridden `{hook_name}` in both"):
+            hook = trainer._data_connector._datahook_selector.get_hook(hook_name)
+
+        warning_cache.clear()
+        assert hook == getattr(dm, hook_name)
+
+    def test_with_datamodule_override_model(self, hook_name):
+        model, dm, trainer = self.reset_instances()
+        trainer._data_connector.attach_datamodule(model, datamodule=dm)
+        setattr(model, hook_name, self.overridden_func)
+        with pytest.warns(UserWarning, match=f"have overridden `{hook_name}` in `LightningModule`"):
+            hook = trainer._data_connector._datahook_selector.get_hook(hook_name)
+
+        warning_cache.clear()
+        assert hook == getattr(model, hook_name)
+
+
+def test_invalid_hook_passed_in_datahook_selector():
+    dh_selector = _DataHookSelector(BoringModel(), None)
+    with pytest.raises(ValueError, match="is not a shared hook"):
+        dh_selector.get_hook("setup")
+
+
 def test_eval_distributed_sampler_warning(tmpdir):
     """Test that a warning is raised when `DistributedSampler` is used with evaluation."""