We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_lightning/trainer/train_loop_mixin.py`
Content:
```
1 import numpy as np
2 
3 try:
4     from apex import amp
5 
6     APEX_AVAILABLE = True
7 except ImportError:
8     APEX_AVAILABLE = False
9 
10 
11 class TrainerTrainLoopMixin(object):
12 
13     def train(self):
14         # run all epochs
15         for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
16             # set seed for distributed sampler (enables shuffling for each epoch)
17             if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):
18                 self.get_train_dataloader().sampler.set_epoch(epoch_nb)
19 
20             # get model
21             model = self.get_model()
22 
23             # update training progress in trainer and model
24             model.current_epoch = epoch_nb
25             self.current_epoch = epoch_nb
26             self.total_batches = self.nb_training_batches + self.nb_val_batches
27             self.batch_loss_value = 0  # accumulated grads
28 
29             # limit the number of batches to 1 in fast_dev_run
30             if self.fast_dev_run:
31                 self.total_batches = 1
32 
33             # init progress_bar when requested
34             if self.show_progress_bar:
35                 nb_iterations = self.total_batches
36 
37                 #  for iterable train loader, the progress bar never ends
38                 if self.is_iterable_train_dataloader:
39                     nb_iterations = float('inf')
40                 self.progress_bar.reset(nb_iterations)
41 
42             # changing gradient according accumulation_scheduler
43             self.accumulation_scheduler.on_epoch_begin(epoch_nb, self)
44 
45             # -----------------
46             # RUN TNG EPOCH
47             # -----------------
48             self.run_training_epoch()
49 
50             # update LR schedulers
51             if self.lr_schedulers is not None:
52                 for lr_scheduler in self.lr_schedulers:
53                     lr_scheduler.step(self.current_epoch)
54 
55             # early stopping
56             met_min_epochs = epoch_nb > self.min_nb_epochs
57             if self.enable_early_stop and (met_min_epochs or self.fast_dev_run):
58                 should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb,
59                                                                     logs=self.callback_metrics)
60                 # stop training
61                 stop = should_stop and met_min_epochs
62                 if stop:
63                     return
64 
65         if self.logger is not None:
66             self.logger.finalize("success")
67 
68     def run_training_epoch(self):
69         # before epoch hook
70         if self.is_function_implemented('on_epoch_start'):
71             model = self.get_model()
72             model.on_epoch_start()
73 
74         # run epoch
75         for batch_nb, batch in enumerate(self.get_train_dataloader()):
76             self.batch_nb = batch_nb
77 
78             model = self.get_model()
79             model.global_step = self.global_step
80 
81             # ---------------
82             # RUN TRAIN STEP
83             # ---------------
84             output = self.run_training_batch(batch, batch_nb)
85             batch_result, grad_norm_dic, batch_step_metrics = output
86 
87             # when returning -1 from train_step, we end epoch early
88             early_stop_epoch = batch_result == -1
89 
90             # ---------------
91             # RUN VAL STEP
92             # ---------------
93             is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
94             can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0
95             should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch)
96 
97             # fast_dev_run always forces val checking after train batch
98             if self.fast_dev_run or should_check_val:
99                 self.run_evaluation(test=self.testing)
100 
101             # when logs should be saved
102             should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch
103             if should_save_log or self.fast_dev_run:
104                 if self.proc_rank == 0 and self.logger is not None:
105                     self.logger.save()
106 
107             # when metrics should be logged
108             should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch
109             if should_log_metrics or self.fast_dev_run:
110                 # logs user requested information to logger
111                 self.log_metrics(batch_step_metrics, grad_norm_dic)
112 
113             self.global_step += 1
114             self.total_batch_nb += 1
115 
116             # end epoch early
117             # stop when the flag is changed or we've gone past the amount
118             # requested in the batches
119             if early_stop_epoch or self.fast_dev_run:
120                 break
121 
122             # stop epoch if we limited nb batches
123             met_batch_limit = batch_nb >= self.nb_training_batches
124             if met_batch_limit:
125                 break
126 
127         # epoch end hook
128         if self.is_function_implemented('on_epoch_end'):
129             model = self.get_model()
130             model.on_epoch_end()
131 
132     def run_training_batch(self, batch, batch_nb):
133         # track grad norms
134         grad_norm_dic = {}
135 
136         # track all metrics for callbacks
137         all_callback_metrics = []
138 
139         # track metrics to log
140         all_log_metrics = []
141 
142         if batch is None:
143             return 0, grad_norm_dic
144 
145         # hook
146         if self.is_function_implemented('on_batch_start'):
147             model_ref = self.get_model()
148             response = model_ref.on_batch_start(batch)
149 
150             if response == -1:
151                 return -1, grad_norm_dic
152 
153         if self.show_progress_bar:
154             self.progress_bar.update(1)
155 
156         # call training_step once per optimizer
157         for opt_idx, optimizer in enumerate(self.optimizers):
158 
159             # wrap the forward step in a closure so second order methods work
160             def optimizer_closure():
161                 # forward pass
162                 output = self.training_forward(batch, batch_nb, opt_idx)
163                 closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output
164 
165                 # track metrics for callbacks
166                 all_callback_metrics.append(callback_metrics)
167 
168                 # track progress bar metrics
169                 self.add_tqdm_metrics(progress_bar_metrics)
170                 all_log_metrics.append(log_metrics)
171 
172                 # accumulate loss
173                 # (if accumulate_grad_batches = 1 no effect)
174                 closure_loss = closure_loss / self.accumulate_grad_batches
175 
176                 # backward pass
177                 if self.use_amp:
178                     with amp.scale_loss(closure_loss, optimizer) as scaled_loss:
179                         scaled_loss.backward()
180                 else:
181                     closure_loss.backward()
182 
183                 # insert after step hook
184                 if self.is_function_implemented('on_after_backward'):
185                     model_ref = self.get_model()
186                     model_ref.on_after_backward()
187 
188                 return closure_loss
189 
190             # calculate loss
191             loss = optimizer_closure()
192 
193             # nan grads
194             if self.print_nan_grads:
195                 self.print_nan_gradients()
196 
197             # track total loss for logging (avoid mem leaks)
198             self.batch_loss_value += loss.item()
199 
200             # gradient update with accumulated gradients
201             if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
202 
203                 # track gradient norms when requested
204                 if batch_nb % self.row_log_interval == 0:
205                     if self.track_grad_norm > 0:
206                         model = self.get_model()
207                         grad_norm_dic = model.grad_norm(self.track_grad_norm)
208 
209                 # clip gradients
210                 self.clip_gradients()
211 
212                 # calls .step(), .zero_grad()
213                 # override function to modify this behavior
214                 model = self.get_model()
215                 model.optimizer_step(self.current_epoch, batch_nb,
216                                      optimizer, opt_idx, optimizer_closure)
217 
218                 # calculate running loss for display
219                 self.running_loss.append(self.batch_loss_value)
220                 self.batch_loss_value = 0
221                 self.avg_loss = np.mean(self.running_loss[-100:])
222 
223                 # update progress bar
224                 if self.show_progress_bar:
225                     # add model specific metrics
226                     tqdm_metrics = self.training_tqdm_dict
227                     self.progress_bar.set_postfix(**tqdm_metrics)
228 
229         # activate batch end hook
230         if self.is_function_implemented('on_batch_end'):
231             model = self.get_model()
232             model.on_batch_end()
233 
234         # collapse all metrics into one dict
235         all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}
236 
237         # track all metrics for callbacks
238         self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()}
239 
240         return 0, grad_norm_dic, all_log_metrics
241 
242     def training_forward(self, batch, batch_nb, opt_idx):
243         """
244         Handle forward for each training case (distributed, single gpu, etc...)
245         :param batch:
246         :param batch_nb:
247         :return:
248         """
249         # ---------------
250         # FORWARD
251         # ---------------
252         # enable not needing to add opt_idx to training_step
253         args = [batch, batch_nb]
254         if len(self.optimizers) > 1:
255             args.append(opt_idx)
256 
257         if self.use_ddp or self.use_ddp2:
258             output = self.model(*args)
259         elif self.use_dp:
260             output = self.model(*args)
261         elif self.single_gpu:
262             gpu_id = 0
263             if type(self.data_parallel_device_ids) is list:
264                 gpu_id = self.data_parallel_device_ids[0]
265             batch = self.transfer_batch_to_gpu(batch, gpu_id)
266             args[0] = batch
267             output = self.model.training_step(*args)
268 
269         else:
270             output = self.model.training_step(*args)
271 
272         # format and reduce outputs accordingly
273         output = self.process_output(output, train=True)
274         loss, progress_bar_metrics, log_metrics, callback_metrics = output
275         return loss, progress_bar_metrics, log_metrics, callback_metrics
276
```

Path: `pytorch_lightning/root_module/hooks.py`
Content:
```
1 import torch
2 
3 
4 class ModelHooks(torch.nn.Module):
5 
6     def on_sanity_check_start(self):
7         """
8         Called before starting evaluate
9         :return:
10         """
11         pass
12 
13     def on_batch_start(self, batch):
14         pass
15 
16     def on_batch_end(self):
17         pass
18 
19     def on_epoch_start(self):
20         pass
21 
22     def on_epoch_end(self):
23         pass
24 
25     def on_pre_performance_check(self):
26         pass
27 
28     def on_post_performance_check(self):
29         pass
30 
31     def on_before_zero_grad(self, optimizer):
32         """
33         Called after optimizer.step() and before optimizer.zero_grad()
34 
35         for optimizer in optimizers:
36             optimizer.step()
37             model.on_before_zero_grad(optimizer) # < ---- called here
38             optimizer.zero_grad
39 
40         :param optimizer:
41         :return:
42         """
43         pass
44 
45     def on_after_backward(self):
46         """
47         Called after loss.backward() and before optimizers do anything
48         :return:
49         """
50         pass
51
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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

{"golden_diff": "diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py\n--- a/pytorch_lightning/root_module/hooks.py\n+++ b/pytorch_lightning/root_module/hooks.py\n@@ -1,6 +1,14 @@\n import torch\n \n \n+try:\n+    from apex import amp\n+\n+    APEX_AVAILABLE = True\n+except ImportError:\n+    APEX_AVAILABLE = False\n+\n+\n class ModelHooks(torch.nn.Module):\n \n     def on_sanity_check_start(self):\n@@ -48,3 +56,17 @@\n         :return:\n         \"\"\"\n         pass\n+\n+    def backward(self, use_amp, loss, optimizer):\n+        \"\"\"\n+        Override backward with your own implementation if you need to\n+        :param use_amp: Whether amp was requested or not\n+        :param loss: Loss is already scaled by accumulated grads\n+        :param optimizer: Current optimizer being used\n+        :return:\n+        \"\"\"\n+        if use_amp:\n+            with amp.scale_loss(loss, optimizer) as scaled_loss:\n+                scaled_loss.backward()\n+        else:\n+            loss.backward()\ndiff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py\n--- a/pytorch_lightning/trainer/train_loop_mixin.py\n+++ b/pytorch_lightning/trainer/train_loop_mixin.py\n@@ -174,11 +174,9 @@\n                 closure_loss = closure_loss / self.accumulate_grad_batches\n \n                 # backward pass\n-                if self.use_amp:\n-                    with amp.scale_loss(closure_loss, optimizer) as scaled_loss:\n-                        scaled_loss.backward()\n-                else:\n-                    closure_loss.backward()\n+                # done in hook so user can overwrite if needed\n+                model_ref = self.get_model()\n+                model_ref.backward(self.use_amp, closure_loss, optimizer)\n \n                 # insert after step hook\n                 if self.is_function_implemented('on_after_backward'):\n", "issue": "Support for retain_graph=True\n**Is your feature request related to a problem? Please describe.**\r\nSome models require retain_graph=True, but it's not possible to set it in the .backward() call inside of Trainer.__run_training_batch(...)\r\n\r\n**Describe the solution you'd like**\r\nAdd train_graph member function the LightningModule have the trainer read this option and then pass it into the .backward() call.\r\n\r\n**Describe alternatives you've considered**\r\nDriving a version of Trainer to support retain_graph=True is tough because __run_training_batch and other functions are name-mangled.\r\n\r\n\n", "before_files": [{"content": "import numpy as np\n\ntry:\n    from apex import amp\n\n    APEX_AVAILABLE = True\nexcept ImportError:\n    APEX_AVAILABLE = False\n\n\nclass TrainerTrainLoopMixin(object):\n\n    def train(self):\n        # run all epochs\n        for epoch_nb in range(self.current_epoch, self.max_nb_epochs):\n            # set seed for distributed sampler (enables shuffling for each epoch)\n            if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):\n                self.get_train_dataloader().sampler.set_epoch(epoch_nb)\n\n            # get model\n            model = self.get_model()\n\n            # update training progress in trainer and model\n            model.current_epoch = epoch_nb\n            self.current_epoch = epoch_nb\n            self.total_batches = self.nb_training_batches + self.nb_val_batches\n            self.batch_loss_value = 0  # accumulated grads\n\n            # limit the number of batches to 1 in fast_dev_run\n            if self.fast_dev_run:\n                self.total_batches = 1\n\n            # init progress_bar when requested\n            if self.show_progress_bar:\n                nb_iterations = self.total_batches\n\n                #  for iterable train loader, the progress bar never ends\n                if self.is_iterable_train_dataloader:\n                    nb_iterations = float('inf')\n                self.progress_bar.reset(nb_iterations)\n\n            # changing gradient according accumulation_scheduler\n            self.accumulation_scheduler.on_epoch_begin(epoch_nb, self)\n\n            # -----------------\n            # RUN TNG EPOCH\n            # -----------------\n            self.run_training_epoch()\n\n            # update LR schedulers\n            if self.lr_schedulers is not None:\n                for lr_scheduler in self.lr_schedulers:\n                    lr_scheduler.step(self.current_epoch)\n\n            # early stopping\n            met_min_epochs = epoch_nb > self.min_nb_epochs\n            if self.enable_early_stop and (met_min_epochs or self.fast_dev_run):\n                should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb,\n                                                                    logs=self.callback_metrics)\n                # stop training\n                stop = should_stop and met_min_epochs\n                if stop:\n                    return\n\n        if self.logger is not None:\n            self.logger.finalize(\"success\")\n\n    def run_training_epoch(self):\n        # before epoch hook\n        if self.is_function_implemented('on_epoch_start'):\n            model = self.get_model()\n            model.on_epoch_start()\n\n        # run epoch\n        for batch_nb, batch in enumerate(self.get_train_dataloader()):\n            self.batch_nb = batch_nb\n\n            model = self.get_model()\n            model.global_step = self.global_step\n\n            # ---------------\n            # RUN TRAIN STEP\n            # ---------------\n            output = self.run_training_batch(batch, batch_nb)\n            batch_result, grad_norm_dic, batch_step_metrics = output\n\n            # when returning -1 from train_step, we end epoch early\n            early_stop_epoch = batch_result == -1\n\n            # ---------------\n            # RUN VAL STEP\n            # ---------------\n            is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0\n            can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0\n            should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch)\n\n            # fast_dev_run always forces val checking after train batch\n            if self.fast_dev_run or should_check_val:\n                self.run_evaluation(test=self.testing)\n\n            # when logs should be saved\n            should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch\n            if should_save_log or self.fast_dev_run:\n                if self.proc_rank == 0 and self.logger is not None:\n                    self.logger.save()\n\n            # when metrics should be logged\n            should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch\n            if should_log_metrics or self.fast_dev_run:\n                # logs user requested information to logger\n                self.log_metrics(batch_step_metrics, grad_norm_dic)\n\n            self.global_step += 1\n            self.total_batch_nb += 1\n\n            # end epoch early\n            # stop when the flag is changed or we've gone past the amount\n            # requested in the batches\n            if early_stop_epoch or self.fast_dev_run:\n                break\n\n            # stop epoch if we limited nb batches\n            met_batch_limit = batch_nb >= self.nb_training_batches\n            if met_batch_limit:\n                break\n\n        # epoch end hook\n        if self.is_function_implemented('on_epoch_end'):\n            model = self.get_model()\n            model.on_epoch_end()\n\n    def run_training_batch(self, batch, batch_nb):\n        # track grad norms\n        grad_norm_dic = {}\n\n        # track all metrics for callbacks\n        all_callback_metrics = []\n\n        # track metrics to log\n        all_log_metrics = []\n\n        if batch is None:\n            return 0, grad_norm_dic\n\n        # hook\n        if self.is_function_implemented('on_batch_start'):\n            model_ref = self.get_model()\n            response = model_ref.on_batch_start(batch)\n\n            if response == -1:\n                return -1, grad_norm_dic\n\n        if self.show_progress_bar:\n            self.progress_bar.update(1)\n\n        # call training_step once per optimizer\n        for opt_idx, optimizer in enumerate(self.optimizers):\n\n            # wrap the forward step in a closure so second order methods work\n            def optimizer_closure():\n                # forward pass\n                output = self.training_forward(batch, batch_nb, opt_idx)\n                closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output\n\n                # track metrics for callbacks\n                all_callback_metrics.append(callback_metrics)\n\n                # track progress bar metrics\n                self.add_tqdm_metrics(progress_bar_metrics)\n                all_log_metrics.append(log_metrics)\n\n                # accumulate loss\n                # (if accumulate_grad_batches = 1 no effect)\n                closure_loss = closure_loss / self.accumulate_grad_batches\n\n                # backward pass\n                if self.use_amp:\n                    with amp.scale_loss(closure_loss, optimizer) as scaled_loss:\n                        scaled_loss.backward()\n                else:\n                    closure_loss.backward()\n\n                # insert after step hook\n                if self.is_function_implemented('on_after_backward'):\n                    model_ref = self.get_model()\n                    model_ref.on_after_backward()\n\n                return closure_loss\n\n            # calculate loss\n            loss = optimizer_closure()\n\n            # nan grads\n            if self.print_nan_grads:\n                self.print_nan_gradients()\n\n            # track total loss for logging (avoid mem leaks)\n            self.batch_loss_value += loss.item()\n\n            # gradient update with accumulated gradients\n            if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:\n\n                # track gradient norms when requested\n                if batch_nb % self.row_log_interval == 0:\n                    if self.track_grad_norm > 0:\n                        model = self.get_model()\n                        grad_norm_dic = model.grad_norm(self.track_grad_norm)\n\n                # clip gradients\n                self.clip_gradients()\n\n                # calls .step(), .zero_grad()\n                # override function to modify this behavior\n                model = self.get_model()\n                model.optimizer_step(self.current_epoch, batch_nb,\n                                     optimizer, opt_idx, optimizer_closure)\n\n                # calculate running loss for display\n                self.running_loss.append(self.batch_loss_value)\n                self.batch_loss_value = 0\n                self.avg_loss = np.mean(self.running_loss[-100:])\n\n                # update progress bar\n                if self.show_progress_bar:\n                    # add model specific metrics\n                    tqdm_metrics = self.training_tqdm_dict\n                    self.progress_bar.set_postfix(**tqdm_metrics)\n\n        # activate batch end hook\n        if self.is_function_implemented('on_batch_end'):\n            model = self.get_model()\n            model.on_batch_end()\n\n        # collapse all metrics into one dict\n        all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}\n\n        # track all metrics for callbacks\n        self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()}\n\n        return 0, grad_norm_dic, all_log_metrics\n\n    def training_forward(self, batch, batch_nb, opt_idx):\n        \"\"\"\n        Handle forward for each training case (distributed, single gpu, etc...)\n        :param batch:\n        :param batch_nb:\n        :return:\n        \"\"\"\n        # ---------------\n        # FORWARD\n        # ---------------\n        # enable not needing to add opt_idx to training_step\n        args = [batch, batch_nb]\n        if len(self.optimizers) > 1:\n            args.append(opt_idx)\n\n        if self.use_ddp or self.use_ddp2:\n            output = self.model(*args)\n        elif self.use_dp:\n            output = self.model(*args)\n        elif self.single_gpu:\n            gpu_id = 0\n            if type(self.data_parallel_device_ids) is list:\n                gpu_id = self.data_parallel_device_ids[0]\n            batch = self.transfer_batch_to_gpu(batch, gpu_id)\n            args[0] = batch\n            output = self.model.training_step(*args)\n\n        else:\n            output = self.model.training_step(*args)\n\n        # format and reduce outputs accordingly\n        output = self.process_output(output, train=True)\n        loss, progress_bar_metrics, log_metrics, callback_metrics = output\n        return loss, progress_bar_metrics, log_metrics, callback_metrics\n", "path": "pytorch_lightning/trainer/train_loop_mixin.py"}, {"content": "import torch\n\n\nclass ModelHooks(torch.nn.Module):\n\n    def on_sanity_check_start(self):\n        \"\"\"\n        Called before starting evaluate\n        :return:\n        \"\"\"\n        pass\n\n    def on_batch_start(self, batch):\n        pass\n\n    def on_batch_end(self):\n        pass\n\n    def on_epoch_start(self):\n        pass\n\n    def on_epoch_end(self):\n        pass\n\n    def on_pre_performance_check(self):\n        pass\n\n    def on_post_performance_check(self):\n        pass\n\n    def on_before_zero_grad(self, optimizer):\n        \"\"\"\n        Called after optimizer.step() and before optimizer.zero_grad()\n\n        for optimizer in optimizers:\n            optimizer.step()\n            model.on_before_zero_grad(optimizer) # < ---- called here\n            optimizer.zero_grad\n\n        :param optimizer:\n        :return:\n        \"\"\"\n        pass\n\n    def on_after_backward(self):\n        \"\"\"\n        Called after loss.backward() and before optimizers do anything\n        :return:\n        \"\"\"\n        pass\n", "path": "pytorch_lightning/root_module/hooks.py"}], "after_files": [{"content": "import numpy as np\n\ntry:\n    from apex import amp\n\n    APEX_AVAILABLE = True\nexcept ImportError:\n    APEX_AVAILABLE = False\n\n\nclass TrainerTrainLoopMixin(object):\n\n    def train(self):\n        # run all epochs\n        for epoch_nb in range(self.current_epoch, self.max_nb_epochs):\n            # set seed for distributed sampler (enables shuffling for each epoch)\n            if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'):\n                self.get_train_dataloader().sampler.set_epoch(epoch_nb)\n\n            # get model\n            model = self.get_model()\n\n            # update training progress in trainer and model\n            model.current_epoch = epoch_nb\n            self.current_epoch = epoch_nb\n            self.total_batches = self.nb_training_batches + self.nb_val_batches\n            self.batch_loss_value = 0  # accumulated grads\n\n            # limit the number of batches to 1 in fast_dev_run\n            if self.fast_dev_run:\n                self.total_batches = 1\n\n            # init progress_bar when requested\n            if self.show_progress_bar:\n                nb_iterations = self.total_batches\n\n                #  for iterable train loader, the progress bar never ends\n                if self.is_iterable_train_dataloader:\n                    nb_iterations = float('inf')\n                self.progress_bar.reset(nb_iterations)\n\n            # changing gradient according accumulation_scheduler\n            self.accumulation_scheduler.on_epoch_begin(epoch_nb, self)\n\n            # -----------------\n            # RUN TNG EPOCH\n            # -----------------\n            self.run_training_epoch()\n\n            # update LR schedulers\n            if self.lr_schedulers is not None:\n                for lr_scheduler in self.lr_schedulers:\n                    lr_scheduler.step(self.current_epoch)\n\n            # early stopping\n            met_min_epochs = epoch_nb > self.min_nb_epochs\n            if self.enable_early_stop and (met_min_epochs or self.fast_dev_run):\n                should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb,\n                                                                    logs=self.callback_metrics)\n                # stop training\n                stop = should_stop and met_min_epochs\n                if stop:\n                    return\n\n        if self.logger is not None:\n            self.logger.finalize(\"success\")\n\n    def run_training_epoch(self):\n        # before epoch hook\n        if self.is_function_implemented('on_epoch_start'):\n            model = self.get_model()\n            model.on_epoch_start()\n\n        # run epoch\n        for batch_nb, batch in enumerate(self.get_train_dataloader()):\n            self.batch_nb = batch_nb\n\n            model = self.get_model()\n            model.global_step = self.global_step\n\n            # ---------------\n            # RUN TRAIN STEP\n            # ---------------\n            output = self.run_training_batch(batch, batch_nb)\n            batch_result, grad_norm_dic, batch_step_metrics = output\n\n            # when returning -1 from train_step, we end epoch early\n            early_stop_epoch = batch_result == -1\n\n            # ---------------\n            # RUN VAL STEP\n            # ---------------\n            is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0\n            can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0\n            should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch)\n\n            # fast_dev_run always forces val checking after train batch\n            if self.fast_dev_run or should_check_val:\n                self.run_evaluation(test=self.testing)\n\n            # when logs should be saved\n            should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch\n            if should_save_log or self.fast_dev_run:\n                if self.proc_rank == 0 and self.logger is not None:\n                    self.logger.save()\n\n            # when metrics should be logged\n            should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch\n            if should_log_metrics or self.fast_dev_run:\n                # logs user requested information to logger\n                self.log_metrics(batch_step_metrics, grad_norm_dic)\n\n            self.global_step += 1\n            self.total_batch_nb += 1\n\n            # end epoch early\n            # stop when the flag is changed or we've gone past the amount\n            # requested in the batches\n            if early_stop_epoch or self.fast_dev_run:\n                break\n\n            # stop epoch if we limited nb batches\n            met_batch_limit = batch_nb >= self.nb_training_batches\n            if met_batch_limit:\n                break\n\n        # epoch end hook\n        if self.is_function_implemented('on_epoch_end'):\n            model = self.get_model()\n            model.on_epoch_end()\n\n    def run_training_batch(self, batch, batch_nb):\n        # track grad norms\n        grad_norm_dic = {}\n\n        # track all metrics for callbacks\n        all_callback_metrics = []\n\n        # track metrics to log\n        all_log_metrics = []\n\n        if batch is None:\n            return 0, grad_norm_dic\n\n        # hook\n        if self.is_function_implemented('on_batch_start'):\n            model_ref = self.get_model()\n            response = model_ref.on_batch_start(batch)\n\n            if response == -1:\n                return -1, grad_norm_dic\n\n        if self.show_progress_bar:\n            self.progress_bar.update(1)\n\n        # call training_step once per optimizer\n        for opt_idx, optimizer in enumerate(self.optimizers):\n\n            # wrap the forward step in a closure so second order methods work\n            def optimizer_closure():\n                # forward pass\n                output = self.training_forward(batch, batch_nb, opt_idx)\n                closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output\n\n                # track metrics for callbacks\n                all_callback_metrics.append(callback_metrics)\n\n                # track progress bar metrics\n                self.add_tqdm_metrics(progress_bar_metrics)\n                all_log_metrics.append(log_metrics)\n\n                # accumulate loss\n                # (if accumulate_grad_batches = 1 no effect)\n                closure_loss = closure_loss / self.accumulate_grad_batches\n\n                # backward pass\n                # done in hook so user can overwrite if needed\n                model_ref = self.get_model()\n                model_ref.backward(self.use_amp, closure_loss, optimizer)\n\n                # insert after step hook\n                if self.is_function_implemented('on_after_backward'):\n                    model_ref = self.get_model()\n                    model_ref.on_after_backward()\n\n                return closure_loss\n\n            # calculate loss\n            loss = optimizer_closure()\n\n            # nan grads\n            if self.print_nan_grads:\n                self.print_nan_gradients()\n\n            # track total loss for logging (avoid mem leaks)\n            self.batch_loss_value += loss.item()\n\n            # gradient update with accumulated gradients\n            if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:\n\n                # track gradient norms when requested\n                if batch_nb % self.row_log_interval == 0:\n                    if self.track_grad_norm > 0:\n                        model = self.get_model()\n                        grad_norm_dic = model.grad_norm(self.track_grad_norm)\n\n                # clip gradients\n                self.clip_gradients()\n\n                # calls .step(), .zero_grad()\n                # override function to modify this behavior\n                model = self.get_model()\n                model.optimizer_step(self.current_epoch, batch_nb,\n                                     optimizer, opt_idx, optimizer_closure)\n\n                # calculate running loss for display\n                self.running_loss.append(self.batch_loss_value)\n                self.batch_loss_value = 0\n                self.avg_loss = np.mean(self.running_loss[-100:])\n\n                # update progress bar\n                if self.show_progress_bar:\n                    # add model specific metrics\n                    tqdm_metrics = self.training_tqdm_dict\n                    self.progress_bar.set_postfix(**tqdm_metrics)\n\n        # activate batch end hook\n        if self.is_function_implemented('on_batch_end'):\n            model = self.get_model()\n            model.on_batch_end()\n\n        # collapse all metrics into one dict\n        all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()}\n\n        # track all metrics for callbacks\n        self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()}\n\n        return 0, grad_norm_dic, all_log_metrics\n\n    def training_forward(self, batch, batch_nb, opt_idx):\n        \"\"\"\n        Handle forward for each training case (distributed, single gpu, etc...)\n        :param batch:\n        :param batch_nb:\n        :return:\n        \"\"\"\n        # ---------------\n        # FORWARD\n        # ---------------\n        # enable not needing to add opt_idx to training_step\n        args = [batch, batch_nb]\n        if len(self.optimizers) > 1:\n            args.append(opt_idx)\n\n        if self.use_ddp or self.use_ddp2:\n            output = self.model(*args)\n        elif self.use_dp:\n            output = self.model(*args)\n        elif self.single_gpu:\n            gpu_id = 0\n            if type(self.data_parallel_device_ids) is list:\n                gpu_id = self.data_parallel_device_ids[0]\n            batch = self.transfer_batch_to_gpu(batch, gpu_id)\n            args[0] = batch\n            output = self.model.training_step(*args)\n\n        else:\n            output = self.model.training_step(*args)\n\n        # format and reduce outputs accordingly\n        output = self.process_output(output, train=True)\n        loss, progress_bar_metrics, log_metrics, callback_metrics = output\n        return loss, progress_bar_metrics, log_metrics, callback_metrics\n", "path": "pytorch_lightning/trainer/train_loop_mixin.py"}, {"content": "import torch\n\n\ntry:\n    from apex import amp\n\n    APEX_AVAILABLE = True\nexcept ImportError:\n    APEX_AVAILABLE = False\n\n\nclass ModelHooks(torch.nn.Module):\n\n    def on_sanity_check_start(self):\n        \"\"\"\n        Called before starting evaluate\n        :return:\n        \"\"\"\n        pass\n\n    def on_batch_start(self, batch):\n        pass\n\n    def on_batch_end(self):\n        pass\n\n    def on_epoch_start(self):\n        pass\n\n    def on_epoch_end(self):\n        pass\n\n    def on_pre_performance_check(self):\n        pass\n\n    def on_post_performance_check(self):\n        pass\n\n    def on_before_zero_grad(self, optimizer):\n        \"\"\"\n        Called after optimizer.step() and before optimizer.zero_grad()\n\n        for optimizer in optimizers:\n            optimizer.step()\n            model.on_before_zero_grad(optimizer) # < ---- called here\n            optimizer.zero_grad\n\n        :param optimizer:\n        :return:\n        \"\"\"\n        pass\n\n    def on_after_backward(self):\n        \"\"\"\n        Called after loss.backward() and before optimizers do anything\n        :return:\n        \"\"\"\n        pass\n\n    def backward(self, use_amp, loss, optimizer):\n        \"\"\"\n        Override backward with your own implementation if you need to\n        :param use_amp: Whether amp was requested or not\n        :param loss: Loss is already scaled by accumulated grads\n        :param optimizer: Current optimizer being used\n        :return:\n        \"\"\"\n        if use_amp:\n            with amp.scale_loss(loss, optimizer) as scaled_loss:\n                scaled_loss.backward()\n        else:\n            loss.backward()\n", "path": "pytorch_lightning/root_module/hooks.py"}]}