update

config/rtdetr_config.py

@@ -97,7 +97,7 @@ class RTDetrBaseConfig(object):
         self.normalize_coords = True
         self.mosaic_prob = 0.0
         self.mixup_prob  = 0.0
-        self.copt_paste  = 0.0
+        self.copy_paste  = 0.0
         self.multi_scale = [0.75, 1.25]
         ## Pixel mean & std
         self.pixel_mean = [123.675, 116.28, 103.53]   # RGB format

config/yolov7_af_config.py

@@ -2,7 +2,7 @@
 
 
 def build_yolov7af_config(args):
-    if args.model == 'yolov7_af_s':
+    if   args.model == 'yolov7_af_s':
         return Yolov7AFSConfig()
     elif args.model == 'yolov7_af_l':
         return Yolov7AFLConfig()

config/yolov8_config.py

@@ -2,7 +2,9 @@
 
 
 def build_yolov8_config(args):
-    if args.model == 'yolov8_s':
+    if   args.model == 'yolov8_n':
+        return Yolov8NConfig()
+    elif args.model == 'yolov8_s':
         return Yolov8SConfig()
     else:
         raise NotImplementedError("No config for model: {}".format(args.model))
@@ -88,8 +90,8 @@ class Yolov8BaseConfig(object):
         self.aug_type = 'yolo'
         self.box_format = 'xyxy'
         self.normalize_coords = False
-        self.mosaic_prob = 1.0
-        self.mixup_prob  = 0.15
+        self.mosaic_prob = 0.0
+        self.mixup_prob  = 0.0
         self.copy_paste  = 0.0           # approximated by the YOLOX's mixup
         self.multi_scale = [0.5, 1.25]   # multi scale: [img_size * 0.5, img_size * 1.25]
         ## Pixel mean & std
@@ -115,6 +117,21 @@ class Yolov8BaseConfig(object):
         for k, v in config_dict.items():
             print("{} : {}".format(k, v))
 
+# YOLOv8-N
+class Yolov8NConfig(Yolov8BaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        # ---------------- Model config ----------------
+        self.width = 0.25
+        self.depth = 0.34
+        self.ratio = 2.0
+        self.scale = "n"
+
+        # ---------------- Data process config ----------------
+        self.mosaic_prob = 1.0
+        self.mixup_prob  = 0.0
+        self.copy_paste  = 0.5
+
 # YOLOv8-S
 class Yolov8SConfig(Yolov8BaseConfig):
     def __init__(self) -> None:
@@ -125,6 +142,51 @@ class Yolov8SConfig(Yolov8BaseConfig):
         self.ratio = 2.0
         self.scale = "s"
 
+        # ---------------- Data process config ----------------
+        self.mosaic_prob = 1.0
+        self.mixup_prob  = 0.0
+        self.copy_paste  = 0.5
+
+# YOLOv8-M
+class Yolov8MConfig(Yolov8BaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        # ---------------- Model config ----------------
+        self.width = 0.75
+        self.depth = 0.67
+        self.ratio = 1.5
+        self.scale = "m"
+
+        # ---------------- Data process config ----------------
+        self.mosaic_prob = 1.0
+        self.mixup_prob  = 0.1
+        self.copy_paste  = 0.5
+
+# YOLOv8-L
+class Yolov8MConfig(Yolov8BaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        # ---------------- Model config ----------------
+        self.width = 1.0
+        self.depth = 1.0
+        self.ratio = 1.0
+        self.scale = "l"
+
+        # ---------------- Data process config ----------------
+        self.mosaic_prob = 1.0
+        self.mixup_prob  = 0.1
+        self.copy_paste  = 0.5
+
+# YOLOv8-X
+class Yolov8MConfig(Yolov8BaseConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        # ---------------- Model config ----------------
+        self.width = 1.25
+        self.depth = 1.0
+        self.ratio = 1.0
+        self.scale = "x"
+
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0
         self.mixup_prob  = 0.1

engine.py

@@ -354,7 +354,7 @@ class RTDetrTrainer(object):
         self.optimizer, self.start_epoch = build_rtdetr_optimizer(cfg, model, args.resume)
 
         # ---------------------------- Build LR Scheduler ----------------------------
-        self.wp_lr_scheduler = LinearWarmUpLrScheduler(cfg.base_lr, wp_iter=cfg.warmup_iters)
+        self.wp_lr_scheduler = LinearWarmUpLrScheduler(cfg.warmup_iters, cfg.base_lr)
         self.lr_scheduler    = build_lr_scheduler(cfg, self.optimizer, args.resume)
 
     def train(self, model):
@@ -386,43 +386,43 @@ class RTDetrTrainer(object):
         # set eval mode
         model.eval()
         model_eval = model if self.model_ema is None else self.model_ema.ema
+        cur_map = -1.
+        to_save = False
 
         if distributed_utils.is_main_process():
-            # check evaluator
             if self.evaluator is None:
                 print('No evaluator ... save model and go on training.')
-                print('Saving state, epoch: {}'.format(self.epoch))
+                to_save = True
                 weight_name = '{}_no_eval.pth'.format(self.args.model)
                 checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                torch.save({'model': model_eval.state_dict(),
-                            'mAP': -1.,
-                            'optimizer': self.optimizer.state_dict(),
-                            'lr_scheduler': self.lr_scheduler.state_dict(),
-                            'epoch': self.epoch,
-                            'args': self.args}, 
-                            checkpoint_path)               
             else:
-                print('eval ...')
-                # evaluate
+                print('Eval ...')
+                # Evaluate
                 with torch.no_grad():
                     self.evaluator.evaluate(model_eval)
 
-                # save model
                 cur_map = self.evaluator.map
                 if cur_map > self.best_map:
                     # update best-map
                     self.best_map = cur_map
-                    # save model
-                    print('Saving state, epoch:', self.epoch)
-                    weight_name = '{}_best.pth'.format(self.args.model)
-                    checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                    torch.save({'model': model_eval.state_dict(),
-                                'mAP': round(self.best_map*100, 1),
-                                'optimizer': self.optimizer.state_dict(),
-                                'lr_scheduler': self.lr_scheduler.state_dict(),
-                                'epoch': self.epoch,
-                                'args': self.args}, 
-                                checkpoint_path)                      
+                    to_save = True
+
+            # Save model
+            if to_save:
+                print('Saving state, epoch:', self.epoch)
+                weight_name = '{}_best.pth'.format(self.args.model)
+                checkpoint_path = os.path.join(self.path_to_save, weight_name)
+                state_dicts = {
+                    'model': model_eval.state_dict(),
+                    'mAP': round(cur_map*100, 1),
+                    'optimizer':  self.optimizer.state_dict(),
+                    'lr_scheduler': self.lr_scheduler.state_dict(),
+                    'epoch': self.epoch,
+                    'args': self.args,
+                    }
+                if self.model_ema is not None:
+                    state_dicts["ema_updates"] = self.model_ema.updates
+                torch.save(state_dicts, checkpoint_path)                      
 
         if self.args.distributed:
             # wait for all processes to synchronize
@@ -506,8 +506,8 @@ class RTDetrTrainer(object):
                     self.model_ema.update(model)
 
             # Update log
-            metric_logger.update(loss=losses.item(), **loss_dict_reduced)
-            metric_logger.update(lr=self.optimizer.param_groups[0]["lr"])
+            metric_logger.update(**loss_dict_reduced)
+            metric_logger.update(lr=self.optimizer.param_groups[2]["lr"])
             metric_logger.update(grad_norm=grad_norm)
             metric_logger.update(size=img_size)