Add the on_before_backward hook

CHANGELOG.md

@@ -131,6 +131,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `on_load_checkpoint` and `on_save_checkpoint` hooks to the `PrecisionPlugin` base class ([#7831](https://github.com/PyTorchLightning/pytorch-lightning/pull/7831))
 
 
+- Added `on_before_backward` hook ([#7865](https://github.com/PyTorchLightning/pytorch-lightning/pull/7865))
+
+
 ### Deprecated
 
 

pytorch_lightning/callbacks/base.py

@@ -19,6 +19,7 @@
 import abc
 from typing import Any, Dict, List, Optional
 
+import torch
 from torch.optim import Optimizer
 
 import pytorch_lightning as pl
@@ -296,6 +297,10 @@ def on_load_checkpoint(
         """
         pass
 
+    def on_before_backward(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', loss: torch.Tensor) -> None:
+        """Called before ``loss.backward()``."""
+        pass
+
     def on_after_backward(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
         """Called after ``loss.backward()`` and before optimizers do anything."""
         pass

pytorch_lightning/callbacks/lambda_function.py

@@ -77,6 +77,7 @@ def __init__(
         on_keyboard_interrupt: Optional[Callable] = None,
         on_save_checkpoint: Optional[Callable] = None,
         on_load_checkpoint: Optional[Callable] = None,
+        on_before_backward: Optional[Callable] = None,
         on_after_backward: Optional[Callable] = None,
         on_before_zero_grad: Optional[Callable] = None,
         on_predict_start: Optional[Callable] = None,

pytorch_lightning/core/lightning.py

@@ -1326,6 +1326,23 @@ def training_step(...):
         self.trainer.train_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
         self._running_manual_backward = False
 
+    def on_before_backward(self, loss: torch.Tensor) -> None:
+        """
+        Override on_before_backward with your own implementation if you need to.
+
+        Args:
+            loss: Loss is already scaled by accumulated grads and possibly scaled by Mixed Precision.
+
+        Called to before backward step.
+
+        Example::
+
+            def on_before_backward(self, loss):
+                print(f"Current Loss: {loss}")
+
+        """
+        pass
+
     def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None:
         """
         Override backward with your own implementation if you need to.

pytorch_lightning/plugins/precision/apex_amp.py

@@ -74,6 +74,9 @@ def backward(
         # enter apex context
         closure_loss = scaled_loss.__enter__()
 
+        # hook
+        model.trainer.call_hook("on_before_backward", closure_loss)
+
         # do backward pass
         # TODO: not entirely sure, why we need this
         if model is not None and isinstance(model, LightningModule):

pytorch_lightning/plugins/precision/deepspeed_precision.py

@@ -69,6 +69,10 @@ def backward(
             )
         # todo: hack around for deepspeed engine to call backward
         deepspeed_engine = model.trainer.model
+
+        # hook
+        model.trainer.call_hook("on_before_backward", closure_loss)
+
         deepspeed_engine.backward(closure_loss, *args, **kwargs)
         # once backward has been applied, release graph
         closure_loss = closure_loss.detach()

pytorch_lightning/plugins/precision/native_amp.py

@@ -59,6 +59,7 @@ def backward(
         """
         closure_loss = self.scaler.scale(closure_loss)
 
+        # call `on_before_backward` hook
         closure_loss = super().backward(model, closure_loss, optimizer, opt_idx, should_accumulate, *args, **kwargs)
 
         # unscale gradient to allow analyze within `on_after_backward`

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -73,6 +73,8 @@ def backward(
         """
         automatic_optimization = model.automatic_optimization
 
+        model.trainer.call_hook("on_before_backward", closure_loss)
+
         # do backward pass
         if automatic_optimization:
             model.backward(closure_loss, optimizer, opt_idx)

pytorch_lightning/trainer/callback_hook.py

@@ -17,6 +17,8 @@
 from inspect import signature
 from typing import Any, Callable, Dict, List, Optional, Type
 
+import torch
+
 from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.utilities import rank_zero_deprecation, rank_zero_warn
@@ -313,6 +315,13 @@ def on_load_checkpoint(self, checkpoint):
                 else:
                     callback.on_load_checkpoint(self, self.lightning_module, state)
 
+    def on_before_backward(self, loss: torch.Tensor) -> None:
+        """
+        Called before loss.backward().
+        """
+        for callback in self.callbacks:
+            callback.on_before_backward(self, self.lightning_module, loss)
+
     def on_after_backward(self):
         """
         Called after loss.backward() and before optimizers do anything.

pytorch_lightning/trainer/connectors/logger_connector/fx_validator.py

@@ -21,6 +21,7 @@ class FxValidator:
     functions: Dict[str, Optional[Dict[str, Tuple[bool]]]] = dict(
         on_before_accelerator_backend_setup=None,
         on_configure_sharded_model=None,
+        on_before_backward=dict(on_step=(False, True), on_epoch=(False, True)),
         on_after_backward=dict(on_step=(False, True), on_epoch=(False, True)),
         on_before_zero_grad=dict(on_step=(False, True), on_epoch=(False, True)),
         on_init_start=None,

pytorch_lightning/trainer/trainer.py

@@ -1197,7 +1197,9 @@ def _call_teardown_hook(self, model: LightningModule) -> None:
     def _reset_result_and_set_fx_name(self, hook_name: str) -> bool:
         # on_before_zero_grad is called within training_step
         # TODO(@carmocca): Result should handle this logic
-        if "batch_start" in hook_name or hook_name in ("on_before_zero_grad", "on_after_backward"):
+        if "batch_start" in hook_name or hook_name in (
+            "on_before_zero_grad", "on_before_backward", "on_after_backward"
+        ):
             return True
         model_ref = self.lightning_module
         if model_ref is not None:

tests/callbacks/test_callbacks.py

@@ -68,18 +68,21 @@ def test_trainer_callback_hook_system_fit(_, tmpdir):
         call.on_batch_start(trainer, model),
         call.on_train_batch_start(trainer, model, ANY, 0, 0),
         call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
+        call.on_before_backward(trainer, model, ANY),
         call.on_after_backward(trainer, model),
         call.on_train_batch_end(trainer, model, ANY, ANY, 0, 0),
         call.on_batch_end(trainer, model),
         call.on_batch_start(trainer, model),
         call.on_train_batch_start(trainer, model, ANY, 1, 0),
         call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
+        call.on_before_backward(trainer, model, ANY),
         call.on_after_backward(trainer, model),
         call.on_train_batch_end(trainer, model, ANY, ANY, 1, 0),
         call.on_batch_end(trainer, model),
         call.on_batch_start(trainer, model),
         call.on_train_batch_start(trainer, model, ANY, 2, 0),
         call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
+        call.on_before_backward(trainer, model, ANY),
         call.on_after_backward(trainer, model),
         call.on_train_batch_end(trainer, model, ANY, ANY, 2, 0),
         call.on_batch_end(trainer, model),

tests/core/test_lightning_module.py

@@ -379,3 +379,20 @@ def assert_device(device: torch.device) -> None:
     assert_device(torch.device("cpu"))
     trainer.predict(model, dataloaders=model.train_dataloader())
     assert_device(torch.device("cpu"))
+
+
+def test_on_before_backward(tmpdir):
+
+    class TestModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.on_before_backward_called = False
+
+        def on_before_backward(self, loss: torch.Tensor) -> None:
+            self.on_before_backward_called = True
+
+    model = TestModel()
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
+    trainer.fit(model)
+    assert model.on_before_backward_called

tests/helpers/datamodules.py

@@ -24,6 +24,10 @@
 if _SKLEARN_AVAILABLE:
     from sklearn.datasets import make_classification, make_regression
     from sklearn.model_selection import train_test_split
+else:
+    make_classification = None
+    make_regression = None
+    train_test_split = None
 
 
 class MNISTDataModule(LightningDataModule):

tests/plugins/test_amp_plugins.py

@@ -19,6 +19,7 @@
 import torch
 
 from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import Callback
 from pytorch_lightning.plugins import ApexMixedPrecisionPlugin, NativeMixedPrecisionPlugin
 from pytorch_lightning.plugins.precision import MixedPrecisionPlugin
 from tests.helpers import BoringModel
@@ -69,18 +70,34 @@ def test_amp_apex_ddp(
     assert isinstance(trainer.precision_plugin, plugin_cls)
 
 
+class CheckOnBeforeBackward(Callback):
+
+    def __init__(self):
+        self.on_before_backward_called = False
+
+    def on_before_backward(self, trainer, pl_module, loss):
+        assert isinstance(loss, torch.Tensor)
+        assert loss.grad_fn is not None
+        self.on_before_backward_called = True
+
+
 class GradientUnscaleBoringModel(BoringModel):
 
     def on_after_backward(self):
         norm = torch.nn.utils.clip_grad_norm_(self.parameters(), 2)
         if not (torch.isinf(norm) or torch.isnan(norm)):
             assert norm.item() < 15.
 
+        cb = [cb for cb in self.trainer.callbacks if isinstance(cb, CheckOnBeforeBackward)]
+        assert len(cb) == 1
+        assert cb[0].on_before_backward_called
+
 
 @RunIf(min_gpus=2, amp_native=True)
 @pytest.mark.parametrize('accum', [1, 2])
 def test_amp_gradient_unscale(tmpdir, accum: int):
     model = GradientUnscaleBoringModel()
+    cb = CheckOnBeforeBackward()
 
     trainer = Trainer(
         max_epochs=2,
@@ -95,13 +112,15 @@ def test_amp_gradient_unscale(tmpdir, accum: int):
         track_grad_norm=2,
         log_every_n_steps=1,
         accumulate_grad_batches=accum,
+        callbacks=[cb]
     )
     trainer.fit(model)
 
 
 @RunIf(min_gpus=2, amp_apex=True, special=True)
 @pytest.mark.parametrize("amp_level", ['O2'])
 def test_amp_apex_ddp_fit(amp_level, tmpdir):
+    cb = CheckOnBeforeBackward()
 
     class CustomBoringModel(BoringModel):
 
@@ -110,6 +129,11 @@ def training_step(self, batch, batch_idx):
             assert self.trainer.precision_plugin._connected
             return super().training_step(batch, batch_idx)
 
+        def on_after_backward(self):
+            cb = [cb for cb in self.trainer.callbacks if isinstance(cb, CheckOnBeforeBackward)]
+            assert len(cb) == 1
+            assert cb[0].on_before_backward_called
+
     trainer = Trainer(
         default_root_dir=tmpdir,
         fast_dev_run=True,
@@ -118,6 +142,7 @@ def training_step(self, batch, batch_idx):
         gpus=2,
         accelerator='ddp',
         plugins=ApexMixedPrecisionPlugin(amp_level=amp_level),
+        callbacks=[cb]
     )
     assert isinstance(trainer.precision_plugin, ApexMixedPrecisionPlugin)
     model = CustomBoringModel()

tests/plugins/test_deepspeed_plugin.py

@@ -573,14 +573,25 @@ def test_deepspeed_multigpu_stage_2_accumulated_grad_batches(tmpdir, cpu_offload
 
     class VerificationCallback(Callback):
 
+        def __init__(self):
+            self.on_train_batch_start_called = False
+            self.on_before_backward_called = False
+
         def on_train_batch_start(
             self, trainer, pl_module: LightningModule, batch: Any, batch_idx: int, dataloader_idx: int
         ) -> None:
             deepspeed_engine = trainer.training_type_plugin.model
             assert trainer.global_step == deepspeed_engine.global_steps
+            self.on_train_batch_start_called = True
+
+        def on_before_backward(self, trainer, pl_module, loss):
+            assert isinstance(loss, torch.Tensor)
+            assert loss.grad_fn is not None
+            self.on_before_backward_called = True
 
     model = ModelParallelClassificationModel()
     dm = ClassifDataModule()
+    verification_callback = VerificationCallback()
     trainer = Trainer(
         default_root_dir=tmpdir,
         progress_bar_refresh_rate=0,
@@ -593,6 +604,8 @@ def on_train_batch_start(
         callbacks=[VerificationCallback()]
     )
     trainer.fit(model, datamodule=dm)
+    assert verification_callback.on_train_batch_start_called
+    assert verification_callback.on_before_backward_called
 
 
 @RunIf(min_gpus=2, deepspeed=True, special=True)

tests/trainer/logging_/test_logger_connector.py

@@ -274,6 +274,7 @@ def test_fx_validator(tmpdir):
     funcs_name = sorted([f for f in dir(Callback) if not f.startswith('_')])
 
     callbacks_func = [
+        'on_before_backward',
         'on_after_backward',
         'on_batch_end',
         'on_batch_start',