🌐 [Add] logging of optimizre(learning rate)

yolo/tools/data_loader.py

@@ -9,7 +9,6 @@ from PIL import Image
 from rich.progress import track
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
-from torch.utils.data.distributed import DistributedSampler
 
 from yolo.config.config import DataConfig, DatasetConfig
 from yolo.tools.data_augmentation import *

yolo/tools/solver.py

@@ -83,7 +83,7 @@ class TrainModel(ValidateModel):
         self.trainer.optimizers[0].next_epoch(len(self.train_loader))
 
     def training_step(self, batch, batch_idx):
-        self.trainer.optimizers[0].next_batch()
+        lr_dict = self.trainer.optimizers[0].next_batch()
         batch_size, images, targets, *_ = batch
         predicts = self(images)
         aux_predicts = self.vec2box(predicts["AUX"])
@@ -97,6 +97,7 @@ class TrainModel(ValidateModel):
             batch_size=batch_size,
             rank_zero_only=True,
         )
+        self.log_dict(lr_dict, prog_bar=False, logger=True, on_epoch=False, rank_zero_only=True)
         return loss * batch_size
 
     def configure_optimizers(self):

yolo/utils/model_utils.py

@@ -65,9 +65,12 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
 
     def next_batch(self):
         self.batch_idx += 1
+        lr_dict = dict()
         for lr_idx, param_group in enumerate(self.param_groups):
             min_lr, max_lr = self.min_lr[lr_idx], self.max_lr[lr_idx]
             param_group["lr"] = min_lr + (self.batch_idx) * (max_lr - min_lr) / self.batch_num
+            lr_dict[f"LR/{lr_idx}"] = param_group["lr"]
+        return lr_dict
 
     optimizer_class.next_batch = next_batch
     optimizer_class.next_epoch = next_epoch