Repository: Lightning-AI/pytorch-lightning

Issue:
Support for retain_graph=True
**Is your feature request related to a problem? Please describe.**
Some models require retain_graph=True, but it's not possible to set it in the .backward() call inside of Trainer.__run_training_batch(...)

**Describe the solution you'd like**
Add train_graph member function the LightningModule have the trainer read this option and then pass it into the .backward() call.

**Describe alternatives you've considered**
Driving a version of Trainer to support retain_graph=True is tough because __run_training_batch and other functions are name-mangled.




Discussion:
why don’t we put a hook for the backward pass and you can override if you need retain_graph?  that keeps it flexible for anything else someone wants to do later
perfect! that keeps with the overall philosophy, which has been very useful btw

> On Oct 12, 2019, at 5:47 AM, William Falcon <notifications@github.com> wrote:
> 
> ﻿
> why don’t we put a hook for the backward pass and you can override if you need retain_graph? that keeps it flexible for anything else someone wants to do later
> 
> —
> You are receiving this because you authored the thread.
> Reply to this email directly, view it on GitHub, or unsubscribe.

@yl1991 @momeara did you guys want to submit a PR for this?

Resolving patch:
diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py
--- 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 @@ def on_after_backward(self):
         :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
--- a/pytorch_lightning/trainer/train_loop_mixin.py
+++ b/pytorch_lightning/trainer/train_loop_mixin.py
@@ -174,11 +174,9 @@ def optimizer_closure():
                 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'):


Test patch: