⚡️ [Update] Optimizer, with more fine LR

yolo/config/task/train.yaml

@@ -22,6 +22,7 @@ optimizer:
     lr: 0.01
     weight_decay: 0.0005
     momentum: 0.937
+    nesterov: true
 
 loss:
   objective:

yolo/tools/solver.py

@@ -87,8 +87,9 @@ class ModelTrainer:
         self.model.train()
         total_loss = defaultdict(lambda: torch.tensor(0.0, device=self.device))
         total_samples = 0
-
+        self.optimizer.next_epoch(len(dataloader))
         for batch_size, images, targets, *_ in dataloader:
+            self.optimizer.next_batch()
             loss_each = self.train_one_batch(images, targets)
 
             for loss_name, loss_val in loss_each.items():

yolo/utils/model_utils.py

@@ -52,11 +52,27 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
     conv_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" not in name]
 
     model_parameters = [
-        {"params": bias_params, "nestrov": True, "momentum": 0.937},
-        {"params": conv_params, "weight_decay": 0.0},
-        {"params": norm_params, "weight_decay": 1e-5},
+        {"params": bias_params, "weight_decay": 0},
+        {"params": conv_params},
+        {"params": norm_params, "weight_decay": 0},
     ]
-    return optimizer_class(model_parameters, **optim_cfg.args)
+
+    def next_epoch(self, batch_num):
+        self.min_lr = self.max_lr
+        self.max_lr = [param["lr"] for param in self.param_groups]
+        self.batch_num = batch_num
+        self.batch_idx = 0
+
+    def next_batch(self):
+        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 + 1) * (max_lr - min_lr) / self.batch_num
+
+    optimizer_class.next_batch = next_batch
+    optimizer_class.next_epoch = next_epoch
+    optimizer = optimizer_class(model_parameters, **optim_cfg.args)
+    optimizer.max_lr = [0.1, 0, 0]
+    return optimizer
 
 
 def create_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LRScheduler:
@@ -69,9 +85,9 @@ def create_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LR
     schedule = scheduler_class(optimizer, **schedule_cfg.args)
     if hasattr(schedule_cfg, "warmup"):
         wepoch = schedule_cfg.warmup.epochs
-        lambda1 = lambda epoch: 0.1 + 0.9 * (epoch / wepoch) if epoch < wepoch else 1
-        lambda2 = lambda epoch: 10 - 9 * (epoch / wepoch) if epoch < wepoch else 1
-        warmup_schedule = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2, lambda1])
+        lambda1 = lambda epoch: (epoch + 1) / wepoch if epoch < wepoch else 1
+        lambda2 = lambda epoch: 10 - 9 * ((epoch + 1) / wepoch) if epoch < wepoch else 1
+        warmup_schedule = LambdaLR(optimizer, lr_lambda=[lambda2, lambda1, lambda1])
         schedule = SequentialLR(optimizer, schedulers=[warmup_schedule, schedule], milestones=[2])
     return schedule