yolov6

config/yolov6_config.py

@@ -81,7 +81,7 @@ class Yolov6BaseConfig(object):
         # ---------------- Lr Scheduler config ----------------
         self.warmup_epoch = 3
         self.lr_scheduler = "cosine"
-        self.max_epoch    = 500
+        self.max_epoch    = 300
         self.eval_epoch   = 10
         self.no_aug_epoch = 20
 

models/yolov6/yolov6_backbone.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov6_basic import RepBlock, RepVGGBlock
+    from .yolov6_basic import RepBlock, RepVGGBlock, RepCSPBlock
 except:
-    from  yolov6_basic import RepBlock, RepVGGBlock
+    from  yolov6_basic import RepBlock, RepVGGBlock, RepCSPBlock
 
 
 # --------------------- Yolov3's Backbone -----------------------
@@ -25,41 +25,41 @@ class Yolov6Backbone(nn.Module):
         self.layer_1 = RepVGGBlock(3, self.feat_dims[0],
                                    kernel_size=3, padding=1, stride=2)
         # P2/4
-        self.layer_2 = nn.Sequential(
-            RepVGGBlock(self.feat_dims[0], self.feat_dims[1],
-                        kernel_size=3, padding=1, stride=2),
-            RepBlock(in_channels  = self.feat_dims[1],
-                     out_channels = self.feat_dims[1],
-                     num_blocks   = round(6*cfg.depth))
-        )
+        self.layer_2 = self.make_block(self.feat_dims[0], self.feat_dims[1], round(6*cfg.depth)) 
         # P3/8
-        self.layer_3 = nn.Sequential(
-            RepVGGBlock(self.feat_dims[1], self.feat_dims[2],
-                        kernel_size=3, padding=1, stride=2),
-            RepBlock(in_channels  = self.feat_dims[2],
-                     out_channels = self.feat_dims[2],
-                     num_blocks   = round(12*cfg.depth))
-        )
+        self.layer_3 = self.make_block(self.feat_dims[1], self.feat_dims[2], round(12*cfg.depth)) 
         # P4/16
-        self.layer_4 = nn.Sequential(
-            RepVGGBlock(self.feat_dims[2], self.feat_dims[3],
-                        kernel_size=3, padding=1, stride=2),
-            RepBlock(in_channels  = self.feat_dims[3],
-                     out_channels = self.feat_dims[3],
-                     num_blocks   = round(18*cfg.depth))
-        )
+        self.layer_4 = self.make_block(self.feat_dims[2], self.feat_dims[3], round(18*cfg.depth)) 
         # P5/32
-        self.layer_5 = nn.Sequential(
-            RepVGGBlock(self.feat_dims[3], self.feat_dims[4],
-                        kernel_size=3, padding=1, stride=2),
-            RepBlock(in_channels  = self.feat_dims[4],
-                     out_channels = self.feat_dims[4],
-                     num_blocks   = round(6*cfg.depth))
-        )
+        self.layer_5 = self.make_block(self.feat_dims[3], self.feat_dims[4], round(6*cfg.depth)) 
 
         # Initialize all layers
         self.init_weights()
-        
+    
+    def make_block(self, in_dim, out_dim, num_blocks=1):
+        if self.model_scale in ["s", "t", "n"]:
+            block = nn.Sequential(
+                RepVGGBlock(in_dim, out_dim,
+                            kernel_size=3, padding=1, stride=2),
+                RepBlock(in_channels  = out_dim,
+                         out_channels = out_dim,
+                         num_blocks   = num_blocks,
+                         block        = RepVGGBlock)
+                         )
+        elif self.model_scale in ["m", "l", "x"]:
+            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)
+                            )
+        else:
+            raise NotImplementedError("Unknown model scale: {}".format(self.model_scale))
+            
+        return block
+
     def init_weights(self):
         """Initialize the parameters."""
         for m in self.modules():

models/yolov6/yolov6_basic.py

@@ -175,81 +175,60 @@ class RepVGGBlock(nn.Module):
 
 
 # ---------------------------- Basic Modules ----------------------------
-class YoloBottleneck(nn.Module):
-    def __init__(self,
-                 in_dim       :int,
-                 out_dim      :int,
-                 kernel_size  :List  = [1, 3],
-                 expansion    :float = 0.5,
-                 shortcut     :bool  = False,
-                 act_type     :str   = 'silu',
-                 norm_type    :str   = 'BN',
-                 depthwise    :bool  = False,
-                 ) -> None:
-        super(YoloBottleneck, self).__init__()
-        inter_dim = int(out_dim * expansion)
-        # ----------------- Network setting -----------------
-        self.conv_layer1 = BasicConv(in_dim, inter_dim,
-                                     kernel_size=kernel_size[0], padding=kernel_size[0]//2, stride=1,
-                                     act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        self.conv_layer2 = BasicConv(inter_dim, out_dim,
-                                     kernel_size=kernel_size[1], padding=kernel_size[1]//2, stride=1,
-                                     act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        self.shortcut = shortcut and in_dim == out_dim
+class RepBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, num_blocks=1, block=RepVGGBlock):
+        super().__init__()
+        self.conv1 = block(in_channels, out_channels)
+        self.block = nn.Sequential(*(block(out_channels, out_channels)
+                                     for _ in range(num_blocks - 1))) if num_blocks > 1 else nn.Identity()
+        if block == BottleRep:
+            self.conv1 = BottleRep(in_channels, out_channels, weight=True)
+            num_blocks = num_blocks // 2
+            self.block = nn.Sequential(*(BottleRep(out_channels, out_channels, weight=True)
+                                         for _ in range(num_blocks - 1))) if num_blocks > 1 else None
 
     def forward(self, x):
-        h = self.conv_layer2(self.conv_layer1(x))
+        x = self.conv1(x)
+        if self.block is not None:
+            x = self.block(x)
 
-        return x + h if self.shortcut else h
+        return x
 
-class CSPBlock(nn.Module):
-    def __init__(self,
-                 in_dim,
-                 out_dim,
-                 num_blocks   :int   = 1,
-                 expansion    :float = 0.5,
-                 shortcut     :bool  = False,
-                 act_type     :str   = 'silu',
-                 norm_type    :str   = 'BN',
-                 depthwise    :bool  = False,
-                 ):
-        super(CSPBlock, self).__init__()
-        # ---------- Basic parameters ----------
-        self.num_blocks = num_blocks
-        self.expansion = expansion
-        self.shortcut = shortcut
-        inter_dim = round(out_dim * expansion)
-        # ---------- Model parameters ----------
-        self.conv_layer_1 = BasicConv(in_dim, inter_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
-        self.conv_layer_2 = BasicConv(in_dim, inter_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
-        self.conv_layer_3 = BasicConv(inter_dim * 2, out_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
-        self.module       = nn.Sequential(*[YoloBottleneck(inter_dim,
-                                                           inter_dim,
-                                                           kernel_size  = [1, 3],
-                                                           expansion    = 1.0,
-                                                           shortcut     = shortcut,
-                                                           act_type     = act_type,
-                                                           norm_type    = norm_type,
-                                                           depthwise    = depthwise)
-                                                           for _ in range(num_blocks)
-                                                           ])
+class BottleRep(nn.Module):
+
+    def __init__(self, in_channels, out_channels, weight=False):
+        super().__init__()
+        self.conv1 = RepVGGBlock(in_channels, out_channels, kernel_size=3, padding=1, stride=1)
+        self.conv2 = RepVGGBlock(out_channels, out_channels, kernel_size=3, padding=1, stride=1)
+        if in_channels != out_channels:
+            self.shortcut = False
+        else:
+            self.shortcut = True
+        if weight:
+            self.alpha = nn.Parameter(torch.ones(1))
+        else:
+            self.alpha = 1.0
 
     def forward(self, x):
-        x1 = self.conv_layer_1(x)
-        x2 = self.module(self.conv_layer_2(x))
-        out = self.conv_layer_3(torch.cat([x1, x2], dim=1))
+        outputs = self.conv1(x)
+        outputs = self.conv2(outputs)
 
-        return out
+        return outputs + self.alpha * x if self.shortcut else outputs
 
-class RepBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, num_blocks=1):
+class RepCSPBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, num_blocks=1, expansion=0.5):
         super().__init__()
-        self.conv1 = RepVGGBlock(in_channels, out_channels, kernel_size=3, padding=1, stride=1)
-        self.block = nn.Sequential(*(RepVGGBlock(out_channels, out_channels, kernel_size=3, padding=1, stride=1)
-                                     for _ in range(num_blocks - 1))) if num_blocks > 1 else nn.Identity()
+        inter_dim = round(out_channels * expansion)  # hidden channels
+        self.cv1 = BasicConv(in_channels, inter_dim, kernel_size=1, act_type='relu')
+        self.cv2 = BasicConv(in_channels, inter_dim, kernel_size=1, act_type='relu')
+        self.cv3 = BasicConv(2 * inter_dim, out_channels, kernel_size=1, act_type='relu')
+
+        self.module = RepBlock(inter_dim, inter_dim, num_blocks, block=BottleRep)
 
     def forward(self, x):
-        x = self.conv1(x)
-        x = self.block(x)
+        x1 = self.cv1(x)
+        x2 = self.module(self.cv2(x))
+        out = self.cv3(torch.cat((x1, x2), dim=1))
+
+        return out
 
-        return x

models/yolov6/yolov6_pafpn.py

@@ -4,9 +4,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 try:
-    from .yolov6_basic import BasicConv, RepBlock
+    from .yolov6_basic import BasicConv, RepBlock, RepCSPBlock
 except:
-    from  yolov6_basic import BasicConv, RepBlock
+    from  yolov6_basic import BasicConv, RepBlock, RepCSPBlock
 
 
 # Yolov6FPN
@@ -14,6 +14,7 @@ class Yolov6PaFPN(nn.Module):
     def __init__(self, cfg, in_dims: List = [256, 512, 1024]):
         super(Yolov6PaFPN, self).__init__()
         self.in_dims = in_dims
+        self.model_scale = cfg.scale
         c3, c4, c5 = in_dims
 
         # ---------------------- Yolov6's Top down FPN ----------------------
@@ -21,33 +22,34 @@ class Yolov6PaFPN(nn.Module):
         self.reduce_layer_1   = BasicConv(c5, round(256*cfg.width),
                                           kernel_size=1, padding=0, stride=1,
                                           act_type=cfg.fpn_act, norm_type=cfg.fpn_norm)
-        self.top_down_layer_1 = RepBlock(in_channels  = c4 + round(256*cfg.width),
-                                         out_channels = round(256*cfg.width),
-                                         num_blocks   = round(12*cfg.depth))
+        self.top_down_layer_1 = self.make_block(in_dim     = c4 + round(256*cfg.width),
+                                                out_dim    = round(256*cfg.width),
+                                                num_blocks = round(12*cfg.depth))
 
         ## P4 -> P3
         self.reduce_layer_2   = BasicConv(round(256*cfg.width), round(128*cfg.width),
                                           kernel_size=1, padding=0, stride=1,
                                           act_type=cfg.fpn_act, norm_type=cfg.fpn_norm)
-        self.top_down_layer_2 = RepBlock(in_channels  = c3 + round(128*cfg.width),
-                                         out_channels = round(128*cfg.width),
-                                         num_blocks   = round(12*cfg.depth))
+        self.top_down_layer_2 = self.make_block(in_dim     = c3 + round(128*cfg.width),
+                                                out_dim    = round(128*cfg.width),
+                                                num_blocks = round(12*cfg.depth))
         
         # ---------------------- Yolov6's Bottom up PAN ----------------------
         ## P3 -> P4
         self.downsample_layer_1 = BasicConv(round(128*cfg.width), round(128*cfg.width),
                                             kernel_size=3, padding=1, stride=2,
                                             act_type=cfg.fpn_act, norm_type=cfg.fpn_norm, depthwise=cfg.fpn_depthwise)
-        self.bottom_up_layer_1  = RepBlock(in_channels  = round(128*cfg.width) + round(128*cfg.width),
-                                           out_channels = round(256*cfg.width),
-                                           num_blocks   = round(12*cfg.depth))
+        self.bottom_up_layer_1  = self.make_block(in_dim     = round(128*cfg.width) + round(128*cfg.width),
+                                                  out_dim    = round(256*cfg.width),
+                                                  num_blocks = round(12*cfg.depth))
+
         ## P4 -> P5
         self.downsample_layer_2 = BasicConv(round(256*cfg.width), round(256*cfg.width),
                                             kernel_size=3, padding=1, stride=2,
                                             act_type=cfg.fpn_act, norm_type=cfg.fpn_norm, depthwise=cfg.fpn_depthwise)
-        self.bottom_up_layer_2  = RepBlock(in_channels  = round(256*cfg.width) + round(256*cfg.width),
-                                           out_channels = round(512*cfg.width),
-                                           num_blocks   = round(12*cfg.depth))
+        self.bottom_up_layer_2  = self.make_block(in_dim     = round(256*cfg.width) + round(256*cfg.width),
+                                                  out_dim    = round(512*cfg.width),
+                                                  num_blocks = round(12*cfg.depth))
 
         # ---------------------- Yolov6's output projection ----------------------
         self.out_layers = nn.ModuleList([
@@ -57,6 +59,20 @@ class Yolov6PaFPN(nn.Module):
                       ])
         self.out_dims = [round(128*cfg.width), round(256*cfg.width), round(512*cfg.width)]
 
+    def make_block(self, in_dim, out_dim, num_blocks=1):
+        if self.model_scale in ["s", "t", "n"]:
+            block = RepBlock(in_channels  = in_dim,
+                             out_channels = out_dim,
+                             num_blocks   = num_blocks)
+        elif self.model_scale in ["m", "l", "x"]:
+            block = RepCSPBlock(in_channels  = in_dim,
+                                out_channels = out_dim,
+                                num_blocks   = num_blocks,
+                                expansion    = 0.5)
+        else:
+            raise NotImplementedError("Unknown model scale: {}".format(self.model_scale))
+            
+        return block        
     def forward(self, features):
         c3, c4, c5 = features
         

train.py

@@ -155,9 +155,6 @@ def train():
     model, criterion = build_model(args, cfg, is_val=True)
     model = model.to(device).train()
     model_without_ddp = model
-    for m in model.modules():
-        if hasattr(m, "switch_to_deploy"):
-            m.switch_to_deploy()
 
     # ---------------------------- Build Model-EMA ----------------------------
     if cfg.use_ema and distributed_utils.get_rank() in [-1, 0]: