add hooks

Author: rohitgr7

pytorch_lightning/accelerators/accelerator.py

@@ -69,7 +69,7 @@ def train_or_test(self):
     def batch_to_device(self, batch: Any, device: torch.device):
         model = self.trainer.get_model()
         if model is not None:
-            return model.transfer_batch_to_device(batch, device)
+            return model.prepare_batch_for_transfer(batch, device)
         return move_data_to_device(batch, device)
 
     def training_step_end(self, output):

pytorch_lightning/core/datamodule.py

@@ -94,7 +94,11 @@ def wrapped_fn(*args, **kwargs):
     return wrapped_fn
 
 
-class LightningDataModule(DataHooks, CheckpointHooks, metaclass=_DataModuleWrapper):
+class LightningDataModule(
+    CheckpointHooks,
+    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.
@@ -247,22 +251,6 @@ def prepare_data(self, *args, **kwargs):
     def setup(self, stage: Optional[str] = None):
         pass
 
-    @abstractmethod
-    def train_dataloader(self, *args, **kwargs) -> DataLoader:
-        pass
-
-    @abstractmethod
-    def val_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
-        pass
-
-    @abstractmethod
-    def test_dataloader(self, *args, **kwargs) -> Union[DataLoader, List[DataLoader]]:
-        pass
-
-    @abstractmethod
-    def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
-        pass
-
     @classmethod
     def add_argparse_args(cls, parent_parser: ArgumentParser) -> ArgumentParser:
         r"""Extends existing argparse by default `LightningDataModule` attributes.

pytorch_lightning/core/decorators.py

@@ -59,8 +59,8 @@ def auto_transfer_args(self, *args, **kwargs):
         if not isinstance(self, LightningModule):
             return fn(self, *args, **kwargs)
 
-        args = self.transfer_batch_to_device(args, self.device)
-        kwargs = self.transfer_batch_to_device(kwargs, self.device)
+        args = self.transfer_batch_to_device(args)
+        kwargs = self.transfer_batch_to_device(kwargs)
         return fn(self, *args, **kwargs)
 
     return auto_transfer_args

pytorch_lightning/core/hooks.py

@@ -14,7 +14,7 @@
 
 """Various hooks to be used in the Lightning code."""
 
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 from pytorch_lightning.utilities import AMPType, move_data_to_device, rank_zero_warn
@@ -300,7 +300,7 @@ def on_after_backward(self):
 
 
 class DataHooks:
-    """Hooks to be used with LightningDataModule."""
+    """Hooks to be used for data related stuff."""
     def prepare_data(self) -> None:
         """
         Use this to download and prepare data.
@@ -508,7 +508,7 @@ def val_dataloader(self):
             will have an argument ``dataloader_idx`` which matches the order here.
         """
 
-    def transfer_batch_to_device(self, batch: Any, device: torch.device) -> Any:
+    def transfer_batch_to_device(self, batch: Any, device: Optional[torch.device] = None) -> Any:
         """
         Override this hook if your :class:`~torch.utils.data.DataLoader` returns tensors
         wrapped in a custom data structure.
@@ -556,8 +556,21 @@ def transfer_batch_to_device(self, batch, device)
             - :func:`~pytorch_lightning.utilities.apply_func.move_data_to_device`
             - :func:`~pytorch_lightning.utilities.apply_func.apply_to_collection`
         """
+        device = device or self.device
         return move_data_to_device(batch, device)
 
+    def on_before_batch_transfer(self, batch):
+        return batch
+
+    def on_after_batch_transfer(self, batch):
+        return batch
+
+    def prepare_batch_for_transfer(self, batch: Any, device: Optional[torch.device] = None):
+        batch = self.on_before_batch_transfer(batch)
+        batch = self.transfer_batch_to_device(batch, device)
+        batch = self.on_after_batch_transfer(batch)
+        return batch
+
 
 class CheckpointHooks:
     """Hooks to be used with Checkpointing."""

pytorch_lightning/core/lightning.py

@@ -55,12 +55,12 @@
 
 class LightningModule(
     ABC,
+    CheckpointHooks,
+    DataHooks,
     DeviceDtypeModuleMixin,
     GradInformation,
-    ModelIO,
     ModelHooks,
-    DataHooks,
-    CheckpointHooks,
+    ModelIO,
     Module,
 ):
     # Below is for property support of JIT in PyTorch 1.7
@@ -1787,7 +1787,7 @@ def to_torchscript(
                 if example_inputs is None:
                     example_inputs = self.example_input_array
                 # automatically send example inputs to the right device and use trace
-                example_inputs = self.transfer_batch_to_device(example_inputs, device=self.device)
+                example_inputs = self.prepare_batch_for_transfer(example_inputs)
                 torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
             else:
                 raise ValueError(f"The 'method' parameter only supports 'script' or 'trace', but value given was:"

pytorch_lightning/core/memory.py

@@ -219,7 +219,7 @@ def _forward_example_input(self) -> None:
         trainer = self._model.trainer
 
         input_ = model.example_input_array
-        input_ = model.transfer_batch_to_device(input_, model.device)
+        input_ = model.prepare_batch_for_transfer(input_, model.device)
 
         if trainer is not None and trainer.amp_backend == AMPType.NATIVE and not trainer.use_tpu:
             model.forward = torch.cuda.amp.autocast()(model.forward)

pytorch_lightning/loggers/tensorboard.py

@@ -198,7 +198,7 @@ def log_graph(self, model: LightningModule, input_array=None):
                 input_array = model.example_input_array
 
             if input_array is not None:
-                input_array = model.transfer_batch_to_device(input_array, model.device)
+                input_array = model.prepare_batch_for_transfer(input_array)
                 self.experiment.add_graph(model, input_array)
             else:
                 rank_zero_warn('Could not log computational graph since the'

pytorch_lightning/loggers/test_tube.py

@@ -146,11 +146,8 @@ def log_graph(self, model: LightningModule, input_array=None):
                 input_array = model.example_input_array
 
             if input_array is not None:
-                self.experiment.add_graph(
-                    model,
-                    model.transfer_batch_to_device(
-                        model.example_input_array, model.device)
-                )
+                input_array = self.prepare_batch_for_transfer(input_array)
+                self.experiment.add_graph(model, input_array)
             else:
                 rank_zero_warn('Could not log computational graph since the'
                                ' `model.example_input_array` attribute is not set'

pytorch_lightning/trainer/connectors/data_connector.py

@@ -118,8 +118,12 @@ def attach_datamodule(self, model, datamodule: Optional[LightningDataModule], st
                 model.test_dataloader = datamodule.test_dataloader
 
             # Override transfer_batch_to_device if dataset-specific to_device logic has been defined in datamodule
+            if is_overridden('on_before_batch_transfer', datamodule):
+                model.on_before_batch_transfer = datamodule.on_before_batch_transfer
             if is_overridden('transfer_batch_to_device', datamodule):
                 model.transfer_batch_to_device = datamodule.transfer_batch_to_device
+            if is_overridden('on_after_batch_transfer', datamodule):
+                model.on_after_batch_transfer = datamodule.on_after_batch_transfer
 
             self.trainer.datamodule = datamodule
             datamodule.trainer = self.trainer