✨ [Update] back the warm up batch step

yolo/tools/solver.py

@@ -27,6 +27,8 @@ class ValidateModel(BaseModel):
         else:
             self.validation_cfg = self.cfg.task.validation
         self.metric = MeanAveragePrecision(iou_type="bbox", box_format="xyxy")
+        self.metric.warn_on_many_detections = False
+        self.val_loader = create_dataloader(self.validation_cfg.data, self.cfg.dataset, self.validation_cfg.task)
 
     def setup(self, stage):
         self.vec2box = create_converter(
@@ -35,7 +37,7 @@ class ValidateModel(BaseModel):
         self.post_proccess = PostProccess(self.vec2box, self.validation_cfg.nms)
 
     def val_dataloader(self):
-        return create_dataloader(self.validation_cfg.data, self.cfg.dataset, self.validation_cfg.task)
+        return self.val_loader
 
     def validation_step(self, batch, batch_idx):
         batch_size, images, targets, rev_tensor, img_paths = batch
@@ -68,15 +70,20 @@ class TrainModel(ValidateModel):
     def __init__(self, cfg: Config):
         super().__init__(cfg)
         self.cfg = cfg
+        self.train_loader = create_dataloader(self.cfg.task.data, self.cfg.dataset, self.cfg.task.task)
 
     def setup(self, stage):
         super().setup(stage)
         self.loss_fn = create_loss_function(self.cfg, self.vec2box)
 
     def train_dataloader(self):
-        return create_dataloader(self.cfg.task.data, self.cfg.dataset, self.cfg.task.task)
+        return self.train_loader
+
+    def on_train_epoch_start(self):
+        self.trainer.optimizers[0].next_epoch(len(self.train_loader))
 
     def training_step(self, batch, batch_idx):
+        self.trainer.optimizers[0].next_batch()
         batch_size, images, targets, *_ = batch
         predicts = self(images)
         aux_predicts = self.vec2box(predicts["AUX"])

yolo/utils/model_utils.py

@@ -56,8 +56,24 @@ def create_optimizer(model: YOLO, optim_cfg: OptimizerConfig) -> Optimizer:
         {"params": conv_params, "momentum": 0.8},
         {"params": norm_params, "momentum": 0.8, "weight_decay": 0},
     ]
+
+    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):
+        self.batch_idx += 1
+        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
+
+    optimizer_class.next_batch = next_batch
+    optimizer_class.next_epoch = next_epoch
+
     optimizer = optimizer_class(model_parameters, **optim_cfg.args)
-    # TODO: implement batch lr schedular when warm up
+    optimizer.max_lr = [0.1, 0, 0]
     return optimizer