chore: log more info to tensorboard (#1608)

chore: log more info to tensorboard
Megvii-BaseDetection · May 15, 2023 · 54a51e1 · 54a51e1
1 parent bb9185c
commit 54a51e1
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/yolox/core/trainer.py b/yolox/core/trainer.py
@@ -264,6 +264,12 @@ def after_iter(self):
             )
 
             if self.rank == 0:
+                if self.args.logger == "tensorboard":
+                    self.tblogger.add_scalar(
+                        "train/lr", self.meter["lr"].latest, self.progress_in_iter)
+                    for k, v in loss_meter.items():
+                        self.tblogger.add_scalar(
+                            f"train/{k}", v.latest, self.progress_in_iter)
                 if self.args.logger == "wandb":
                     metrics = {"train/" + k: v.latest for k, v in loss_meter.items()}
                     metrics.update({