modify yolov6

yolo/config/yolov6_config.py

@@ -10,8 +10,6 @@ def build_yolov6_config(args):
         return Yolov6MConfig()
     elif args.model == 'yolov6_l':
         return Yolov6LConfig()
-    elif args.model == 'yolov6_x':
-        return Yolov6XConfig()
     else:
         raise NotImplementedError("No config for model: {}".format(args.model))
     
@@ -152,6 +150,7 @@ class Yolov6MConfig(Yolov6BaseConfig):
         self.width = 0.75
         self.depth = 0.67
         self.scale = "m"
+        self.bk_csp_expansion = 0.67
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0
@@ -166,20 +165,7 @@ class Yolov6LConfig(Yolov6BaseConfig):
         self.width = 1.0
         self.depth = 1.0
         self.scale = "l"
-
-        # ---------------- Data process config ----------------
-        self.mosaic_prob = 1.0
-        self.mixup_prob  = 0.1
-        self.copy_paste  = 0.5
-
-# YOLOv6-X
-class Yolov6XConfig(Yolov6BaseConfig):
-    def __init__(self) -> None:
-        super().__init__()
-        # ---------------- Model config ----------------
-        self.width = 1.25
-        self.depth = 1.34
-        self.scale = "x"
+        self.bk_csp_expansion = 0.5
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0

yolo/models/yolov6/yolov6.py

@@ -43,8 +43,6 @@ class Yolov6(nn.Module):
         ## Pred
         self.pred     = Yolov6DetPredLayer(cfg, self.fpn.out_dims)
 
-        self.switch_deploy()
-
     def switch_deploy(self,):
         for m in self.modules():
             if hasattr(m, "switch_to_deploy"):

yolo/models/yolov6/yolov6_backbone.py

@@ -13,6 +13,7 @@ class Yolov6Backbone(nn.Module):
     def __init__(self, cfg):
         super(Yolov6Backbone, self).__init__()
         # ------------------ Basic setting ------------------
+        self.cfg = cfg
         self.model_scale = cfg.scale
         self.feat_dims = [round(64   * cfg.width),
                           round(128  * cfg.width),
@@ -46,14 +47,14 @@ class Yolov6Backbone(nn.Module):
                          num_blocks   = num_blocks,
                          block        = RepVGGBlock)
                          )
-        elif self.model_scale in ["m", "l", "x"]:
+        elif self.model_scale in ["m", "l"]:
             block = nn.Sequential(
                 RepVGGBlock(in_dim, out_dim,
                             kernel_size=3, padding=1, stride=2),
                 RepCSPBlock(in_channels  = out_dim,
                             out_channels = out_dim,
                             num_blocks   = num_blocks,
-                            expansion    = 0.5)
+                            expansion    = self.cfg.bk_csp_expansion)
                             )
         else:
             raise NotImplementedError("Unknown model scale: {}".format(self.model_scale))