Back hook (#424)

* Fixes #356

* Fixes #356

* Fixes #356

* Fixes #356

* Fixes #356

* Fixes #356

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()

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()

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'):