🔀 [Merge] branch 'TRAIN' into TEST

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/loss_functions.py

@@ -88,7 +88,7 @@ class YOLOLoss:
     def __call__(self, predicts: List[Tensor], targets: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
         predicts_cls, predicts_anc, predicts_box = predicts
         # For each predicted targets, assign a best suitable ground truth box.
-        align_targets, valid_masks = self.matcher(targets, (predicts_cls, predicts_box))
+        align_targets, valid_masks = self.matcher(targets, (predicts_cls.detach(), predicts_box.detach()))
 
         targets_cls, targets_bbox = self.separate_anchor(align_targets)
         predicts_box = predicts_box / self.vec2box.scaler[None, :, None]
@@ -134,6 +134,7 @@ class DualLoss:
 
 
 def create_loss_function(cfg: Config, vec2box) -> DualLoss:
+    # TODO: make it flexible, if cfg doesn't contain aux, only use SingleLoss
     loss_function = DualLoss(cfg, vec2box)
     logger.info("✅ Success load loss function")
     return loss_function

yolo/tools/solver.py

@@ -78,6 +78,8 @@ class ModelTrainer:
             loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
 
         self.scaler.scale(loss).backward()
+        self.scaler.unscale_(self.optimizer)
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=10.0)
         self.scaler.step(self.optimizer)
         self.scaler.update()
 
@@ -87,8 +89,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/logging_utils.py

@@ -67,6 +67,7 @@ def set_seed(seed):
 
 class ProgressLogger(Progress):
     def __init__(self, cfg: Config, exp_name: str, *args, **kwargs):
+        set_seed(cfg.lucky_number)
         local_rank = int(os.getenv("LOCAL_RANK", "0"))
         self.quite_mode = local_rank or getattr(cfg, "quite", False)
         custom_logger(self.quite_mode)

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