Fix broken gradients logging and add lr logging to tensorboard (facebookresearch#1158)

Summary:
Pull Request resolved: facebookresearch#1158

This should help to monitor lr when using warmup/annealing etc

Reviewed By: geof90

Differential Revision: D18624642

fbshipit-source-id: 0cf55150f40c8a3ddf459d9d968f15f58356c488

Author: arbabu123

pytext/metric_reporters/channel.py

@@ -171,6 +171,7 @@ def report(
         meta,
         model,
         *args,
+        optimizer=None,
     ):
         """
         Defines how to format and report data to TensorBoard using the summary
@@ -213,17 +214,19 @@ def report(
                 self.add_scalars(prefix, metrics, epoch)
 
         if stage == Stage.TRAIN:
+            if optimizer is not None:
+                for idx, param_group in enumerate(optimizer.param_groups):
+                    self.summary_writer.add_scalar(
+                        f"optimizer.lr.param_group.{idx}", param_group["lr"], epoch
+                    )
             for key, val in model.named_parameters():
                 if val is not None and len(val) > 0 and not (val == 0).all():
                     limit = 9.9e19
+                    grad = val.grad
                     val = torch.clamp(val.float(), -limit, limit)
                     self.summary_writer.add_histogram(key, val, epoch)
-                    if (
-                        val.grad is not None
-                        and len(val.grad) > 0
-                        and not (val.grad == 0).all()
-                    ):
-                        grad = torch.clamp(val.grad.float(), -limit, limit)
+                    if grad is not None and len(grad) > 0 and not (grad == 0).all():
+                        grad = torch.clamp(grad.float(), -limit, limit)
                         self.summary_writer.add_histogram(
                             key + "_gradients", grad, epoch
                         )

pytext/metric_reporters/metric_reporter.py

@@ -206,7 +206,9 @@ def get_meta(self):
         """
         return {}
 
-    def report_metric(self, model, stage, epoch, reset=True, print_to_channels=True):
+    def report_metric(
+        self, model, stage, epoch, reset=True, print_to_channels=True, optimizer=None
+    ):
         """
         Calculate metrics and average loss, report all statistic data to channels
 
@@ -241,6 +243,7 @@ def report_metric(self, model, stage, epoch, reset=True, print_to_channels=True)
                         self.all_context,
                         self.get_meta(),
                         model,
+                        optimizer,
                     )
 
         if reset:

pytext/trainers/trainer.py

@@ -509,7 +509,13 @@ def run_epoch(
         if report_metric:
             with timing.time("report metrics"):
                 metrics = metric_reporter.report_metric(
-                    model, state.stage, state.epoch, print_to_channels=(state.rank == 0)
+                    model,
+                    state.stage,
+                    state.epoch,
+                    print_to_channels=(state.rank == 0),
+                    optimizer=getattr(
+                        state, "optimizer", None
+                    ),  # optimizer is not present during test
                 )
         else:
             metric_reporter._reset()