🔨 [Add] Warm up in training proccess

yolo/config/config.py

@@ -53,6 +53,7 @@ class SchedulerArgs:
 class SchedulerConfig:
     type: str
     args: SchedulerArgs
+    warmup: Dict[str, Union[str, int, float]]
 
 
 @dataclass

yolo/config/hyper/default.yaml

@@ -15,12 +15,13 @@ data:
   shuffle: True
   pin_memory: True
 train:
-  epoch: 10
+  epoch: 500
   optimizer:
-    type: Adam
+    type: SGD
     args:
-      lr: 0.001
-      weight_decay: 0.0001
+      lr: 0.01
+      weight_decay: 0.0005
+      momentum: 0.937
   loss:
     objective:
       BCELoss: 0.5
@@ -35,10 +36,13 @@ train:
         iou: 6.0
         cls: 0.5
   scheduler:
-    type: StepLR
+    type: LinearLR
+    warmup:
+      epochs: 3.0
     args:
-      step_size: 10
-      gamma: 0.1
+      total_iters: ${hyper.train.epoch}
+      start_factor: 1
+      end_factor: 0.01
   ema:
       enabled: true
       decay: 0.995

yolo/tools/model_helper.py

@@ -2,9 +2,10 @@ from typing import Any, Dict, Type
 
 import torch
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
+from torch.optim.lr_scheduler import LambdaLR, SequentialLR, _LRScheduler
 
 from yolo.config.config import OptimizerConfig, SchedulerConfig
+from yolo.model.yolo import YOLO
 
 
 class EMA:
@@ -31,21 +32,38 @@ class EMA:
             self.shadow[name].copy_(param.data)
 
 
-def get_optimizer(model_parameters, optim_cfg: OptimizerConfig) -> Optimizer:
+def get_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
     """Create an optimizer for the given model parameters based on the configuration.
 
     Returns:
         An instance of the optimizer configured according to the provided settings.
     """
     optimizer_class: Type[Optimizer] = getattr(torch.optim, optim_cfg.type)
+
+    bias_params = [p for name, p in model.named_parameters() if "bias" in name]
+    norm_params = [p for name, p in model.named_parameters() if "weight" in name and "bn" in name]
+    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},
+    ]
     return optimizer_class(model_parameters, **optim_cfg.args)
 
 
-def get_scheduler(optimizer: Optimizer, schedul_cfg: SchedulerConfig) -> _LRScheduler:
+def get_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LRScheduler:
     """Create a learning rate scheduler for the given optimizer based on the configuration.
 
     Returns:
         An instance of the scheduler configured according to the provided settings.
     """
-    scheduler_class: Type[_LRScheduler] = getattr(torch.optim.lr_scheduler, schedul_cfg.type)
-    return scheduler_class(optimizer, **schedul_cfg.args)
+    scheduler_class: Type[_LRScheduler] = getattr(torch.optim.lr_scheduler, schedule_cfg.type)
+    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 + 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=[lambda1, lambda2, lambda1])
+        schedule = SequentialLR(optimizer, schedulers=[warmup_schedule, schedule], milestones=[2])
+    return schedule

yolo/tools/trainer.py

@@ -19,7 +19,7 @@ class Trainer:
 
         self.model = model.to(device)
         self.device = device
-        self.optimizer = get_optimizer(model.parameters(), train_cfg.optimizer)
+        self.optimizer = get_optimizer(model, train_cfg.optimizer)
         self.scheduler = get_scheduler(self.optimizer, train_cfg.scheduler)
         self.loss_fn = get_loss_function(cfg)
         self.progress = CustomProgress(cfg, save_path, use_wandb=True)
@@ -47,7 +47,6 @@ class Trainer:
     def train_one_epoch(self, dataloader):
         self.model.train()
         total_loss = 0
-        self.progress.start_batch(len(dataloader))
 
         for data, targets in dataloader:
             loss, loss_each = self.train_one_batch(data, targets)
@@ -58,7 +57,6 @@ class Trainer:
         if self.scheduler:
             self.scheduler.step()
 
-        self.progress.finish_batch()
         return total_loss / len(dataloader)
 
     def save_checkpoint(self, epoch: int, filename="checkpoint.pt"):
@@ -80,8 +78,9 @@ class Trainer:
             self.progress.start_train(num_epochs)
             for epoch in range(num_epochs):
 
+                self.progress.start_one_epoch(len(dataloader), self.optimizer, epoch)
                 epoch_loss = self.train_one_epoch(dataloader)
-                self.progress.one_epoch()
+                self.progress.finish_one_epoch()
 
                 logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")
                 if (epoch + 1) % 5 == 0: