From d5ca464cc645a63c3dd92548785cb8e5a7222caf Mon Sep 17 00:00:00 2001
From: William Falcon <waf2107@columbia.edu>
Date: Thu, 24 Oct 2019 07:56:56 -0400
Subject: [PATCH] Back hook (#424)

* Fixes #356

* Fixes #356

* Fixes #356

* Fixes #356

* Fixes #356

* Fixes #356
---
 docs/Trainer/hooks.md                         | 22 +++++++++++++++++++
 pytorch_lightning/root_module/hooks.py        | 22 +++++++++++++++++++
 pytorch_lightning/trainer/train_loop_mixin.py |  8 +++----
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/docs/Trainer/hooks.md b/docs/Trainer/hooks.md
index 6fd8df3d..98ef6b8b 100644
--- a/docs/Trainer/hooks.md
+++ b/docs/Trainer/hooks.md
@@ -115,6 +115,28 @@ def on_before_zero_grad(self, optimizer):
     # do something with the optimizer or inspect it. 
 ```
 
+---
+#### backward
+Called to perform backward step.
+Feel free to override as needed.
+
+The loss passed in has already been scaled for accumulated gradients if requested.
+```python
+def backward(self, use_amp, loss, optimizer):
+    """
+    Override backward with your own implementation if you need to
+    :param use_amp: Whether amp was requested or not
+    :param loss: Loss is already scaled by accumulated grads
+    :param optimizer: Current optimizer being used
+    :return:
+    """
+    if use_amp:
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+    else:
+        loss.backward()
+```
+
 ---
 #### on_after_backward
 Called in the training loop after model.backward()
diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py
index 37d41eb5..580d06e3 100644
--- a/pytorch_lightning/root_module/hooks.py
+++ b/pytorch_lightning/root_module/hooks.py
@@ -1,6 +1,14 @@
 import torch
 
 
+try:
+    from apex import amp
+
+    APEX_AVAILABLE = True
+except ImportError:
+    APEX_AVAILABLE = False
+
+
 class ModelHooks(torch.nn.Module):
 
     def on_sanity_check_start(self):
@@ -48,3 +56,17 @@ class ModelHooks(torch.nn.Module):
         :return:
         """
         pass
+
+    def backward(self, use_amp, loss, optimizer):
+        """
+        Override backward with your own implementation if you need to
+        :param use_amp: Whether amp was requested or not
+        :param loss: Loss is already scaled by accumulated grads
+        :param optimizer: Current optimizer being used
+        :return:
+        """
+        if use_amp:
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+        else:
+            loss.backward()
diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py
index 1dad28ae..8cebe7a1 100644
--- a/pytorch_lightning/trainer/train_loop_mixin.py
+++ b/pytorch_lightning/trainer/train_loop_mixin.py
@@ -174,11 +174,9 @@ class TrainerTrainLoopMixin(object):
                 closure_loss = closure_loss / self.accumulate_grad_batches
 
                 # backward pass
-                if self.use_amp:
-                    with amp.scale_loss(closure_loss, optimizer) as scaled_loss:
-                        scaled_loss.backward()
-                else:
-                    closure_loss.backward()
+                # done in hook so user can overwrite if needed
+                model_ref = self.get_model()
+                model_ref.backward(self.use_amp, closure_loss, optimizer)
 
                 # insert after step hook
                 if self.is_function_implemented('on_after_backward'):