Log lr for param groups (#159)

ultralytics/yolo/engine/trainer.py

@@ -317,7 +317,8 @@ class BaseTrainer:
 
                 self.run_callbacks("on_train_batch_end")
 
-            lr = {f"lr{ir}": x['lr'] for ir, x in enumerate(self.optimizer.param_groups)}  # for loggers
+            self.lr = {f"lr/pg{ir}": x['lr'] for ir, x in enumerate(self.optimizer.param_groups)}  # for loggers
+
             self.scheduler.step()
             self.run_callbacks("on_train_epoch_end")
 
@@ -328,7 +329,7 @@ class BaseTrainer:
                 final_epoch = (epoch + 1 == self.epochs)
                 if self.args.val or final_epoch:
                     self.metrics, self.fitness = self.validate()
-                self.save_metrics(metrics={**self.label_loss_items(self.tloss), **self.metrics, **lr})
+                self.save_metrics(metrics={**self.label_loss_items(self.tloss), **self.metrics, **self.lr})
 
                 # Save model
                 if self.args.save or (epoch + 1 == self.epochs):

ultralytics/yolo/utils/callbacks/wb.py

@@ -25,6 +25,7 @@ def on_fit_epoch_end(trainer):
 
 def on_train_epoch_end(trainer):
     wandb.run.log(trainer.label_loss_items(trainer.tloss, prefix="train"), step=trainer.epoch + 1)
+    wandb.run.log(trainer.lr, step=trainer.epoch + 1)
     if trainer.epoch == 1:
         wandb.run.log({f.stem: wandb.Image(str(f))
                        for f in trainer.save_dir.glob('train_batch*.jpg')},