Repository: Lightning-AI/pytorch-lightning

Issue:
Trainer track_grad_norm always results in 0
**Describe the bug**
The Trainer has a flag `track_grad_norm` which allows us to log the gradient norms to Tensorboard. This flag is checked in the `run_tng_epoch` function after the training step and validation step. However, the training step (`__run_tng_batch`) calls `model.optimizer_step()`, which, in the default implementation, calls `optimizer.zero_grad()`. This results in the tracked gradient norms to be always zero. 

Moreover, there is an optional `__run_evaluation` call in the validatin_step. This results in a call to `model.zero_grad()`, which I assume will also result in zero gradient norms. 



**To Reproduce**
Steps to reproduce the behavior:
1. Train a model with the `track_grad_norm` flag set to True and Tensorboard logging enabled
2. Go to Tensorboard
3. Check the gradient_norms

**Expected behavior**
The gradient norms should not always be zero.


**Desktop (please complete the following information):**
 - OS: Windows
 - Version 0.4.9



Discussion:
I temporarily worked around this by overwriting the `optimizer_step()` function of my LightningModule to not set the gradients to zero, and moved that to the `on_batch_start` hook. However, it's clear that the validation runs also cause the gradients to be zero. See the image attached, where there is one validation batch every 5 training batches. The gradient clearly periodically drops to zero. 
![image](https://user-images.githubusercontent.com/7951058/65946306-29132c00-e436-11e9-8c91-2522c4df7212.png)


good catch! 

1. evaluation does not calculate gradients, so the norms should be zero. the real solution is not to track grad norm during evaluation.
2. [ these lines](https://github.com/williamFalcon/pytorch-lightning/blob/480eed5cb61cb67037cac8b8eb6df9e0fe9286d6/pytorch_lightning/trainer/trainer.py#L1024-L1028), should be moved right before optimizer_step ([this line](https://github.com/williamFalcon/pytorch-lightning/blob/480eed5cb61cb67037cac8b8eb6df9e0fe9286d6/pytorch_lightning/trainer/trainer.py#L1230)).

Would love a PR to fix if you have some time!

Resolving patch:
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -991,7 +991,7 @@ def run_training_epoch(self):
             # ---------------
             # RUN TRAIN STEP
             # ---------------
-            batch_result = self.__run_training_batch(batch, batch_nb)
+            batch_result, grad_norm_dic = self.__run_training_batch(batch, batch_nb)
             early_stop_epoch = batch_result == -1
 
             # ---------------
@@ -1022,10 +1022,7 @@ def run_training_epoch(self):
                     metrics.update(mem_map)
 
                 # add norms
-                if self.track_grad_norm > 0:
-                    model = self.__get_model()
-                    grad_norm_dic = model.grad_norm(self.track_grad_norm)
-                    metrics.update(grad_norm_dic)
+                metrics.update(grad_norm_dic)
 
                 if self.__is_function_implemented('on_training_metrics'):
                     model.on_training_metrics(metrics)
@@ -1177,8 +1174,11 @@ def __print_nan_grads(self):
                 print(param, param.grad)
 
     def __run_training_batch(self, batch, batch_nb):
+        # track grad norms
+        grad_norm_dic = {}
+
         if batch is None:
-            return 0
+            return 0, grad_norm_dic
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
@@ -1186,7 +1186,7 @@ def __run_training_batch(self, batch, batch_nb):
             response = model_ref.on_batch_start(batch)
 
             if response == -1:
-                return -1
+                return -1, grad_norm_dic
 
         if self.show_progress_bar:
             self.progress_bar.update(1)
@@ -1225,6 +1225,13 @@ def __run_training_batch(self, batch, batch_nb):
 
             # gradient update with accumulated gradients
             if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
+
+                # track gradient norms when requested
+                if batch_nb % self.row_log_interval == 0:
+                    if self.track_grad_norm > 0:
+                        model = self.__get_model()
+                        grad_norm_dic = model.grad_norm(self.track_grad_norm)
+
                 # clip gradients
                 self.__clip_gradients()
 
@@ -1249,7 +1256,7 @@ def __run_training_batch(self, batch, batch_nb):
             model = self.__get_model()
             model.on_batch_end()
 
-        return 0
+        return 0, grad_norm_dic
 
     def __run_evaluation(self, test=False):
         # when testing make sure user defined a test step


Test patch: