modify yolov3

yolo/config/yolov3_config.py

@@ -24,26 +24,10 @@ class Yolov3BaseConfig(object):
         self.out_stride = [8, 16, 32]
         self.max_stride = 32
         self.num_levels = 3
-        self.scale      = "b"
+        self.model_scale = "l"
         ## Backbone
-        self.bk_act   = 'silu'
-        self.bk_norm  = 'BN'
-        self.bk_depthwise = False
         self.use_pretrained = True
-        ## Neck
-        self.neck_act       = 'silu'
-        self.neck_norm      = 'BN'
-        self.neck_depthwise = False
-        self.neck_expand_ratio = 0.5
-        self.spp_pooling_size  = 5
-        ## FPN
-        self.fpn_act  = 'silu'
-        self.fpn_norm = 'BN'
-        self.fpn_depthwise = False
         ## Head
-        self.head_act  = 'silu'
-        self.head_norm = 'BN'
-        self.head_depthwise = False
         self.head_dim       = 256
         self.num_cls_head   = 2
         self.num_reg_head   = 2
@@ -127,7 +111,7 @@ class Yolov3NConfig(Yolov3BaseConfig):
         # ---------------- Model config ----------------
         self.width = 0.25
         self.depth = 0.34
-        self.scale = "n"
+        self.model_scale = "n"
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0
@@ -141,7 +125,7 @@ class Yolov3SConfig(Yolov3BaseConfig):
         # ---------------- Model config ----------------
         self.width = 0.50
         self.depth = 0.34
-        self.scale = "s"
+        self.model_scale = "s"
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0
@@ -155,7 +139,7 @@ class Yolov3MConfig(Yolov3BaseConfig):
         # ---------------- Model config ----------------
         self.width = 0.75
         self.depth = 0.67
-        self.scale = "m"
+        self.model_scale = "m"
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0
@@ -169,7 +153,7 @@ class Yolov3LConfig(Yolov3BaseConfig):
         # ---------------- Model config ----------------
         self.width = 1.0
         self.depth = 1.0
-        self.scale = "l"
+        self.model_scale = "l"
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0
@@ -183,7 +167,7 @@ class Yolov3XConfig(Yolov3BaseConfig):
         # ---------------- Model config ----------------
         self.width = 1.25
         self.depth = 1.34
-        self.scale = "x"
+        self.model_scale = "x"
 
         # ---------------- Data process config ----------------
         self.mosaic_prob = 1.0

yolo/models/yolov3/modules.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+from typing import List
+
+
+# --------------------- Basic modules ---------------------
+class ConvModule(nn.Module):
+    def __init__(self, 
+                 in_dim,        # in channels
+                 out_dim,       # out channels 
+                 kernel_size=1, # kernel size 
+                 padding=0,     # padding
+                 stride=1,      # padding
+                 dilation=1,    # dilation
+                ):
+        super(ConvModule, self).__init__()
+        self.conv = nn.Conv2d(in_dim, out_dim, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=False)
+        self.norm = nn.BatchNorm2d(out_dim)
+        self.act  = nn.SiLU(inplace=True)
+
+    def forward(self, x):
+        return self.act(self.norm(self.conv(x)))
+
+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,
+                 ) -> None:
+        super(YoloBottleneck, self).__init__()
+        inter_dim = int(out_dim * expansion)
+        # ----------------- Network setting -----------------
+        self.conv_layer1 = ConvModule(in_dim, inter_dim, kernel_size=kernel_size[0], padding=kernel_size[0]//2, stride=1)
+        self.conv_layer2 = ConvModule(inter_dim, out_dim, kernel_size=kernel_size[1], padding=kernel_size[1]//2, stride=1)
+        self.shortcut = shortcut and in_dim == out_dim
+
+    def forward(self, x):
+        h = self.conv_layer2(self.conv_layer1(x))
+
+        return x + h if self.shortcut else h
+
+class ResBlock(nn.Module):
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 num_blocks :int   = 1,
+                 expansion  :float = 0.5,
+                 shortcut   :bool  = False,
+                 ):
+        super(ResBlock, self).__init__()
+        # ---------- Basic parameters ----------
+        self.num_blocks = num_blocks
+        self.expansion = expansion
+        self.shortcut = shortcut
+        # ---------- Model parameters ----------
+        module = []
+        for i in range(num_blocks):
+            if i == 0:
+                module.append(YoloBottleneck(in_dim       = in_dim,
+                                             out_dim      = out_dim,
+                                             kernel_size  = [1, 3],
+                                             expansion    = expansion,
+                                             shortcut     = shortcut,
+                                             ))
+            else:
+                module.append(YoloBottleneck(in_dim       = out_dim,
+                                             out_dim      = out_dim,
+                                             kernel_size  = [1, 3],
+                                             expansion    = expansion,
+                                             shortcut     = shortcut,
+                                             ))
+
+
+        self.module = nn.Sequential(*module)
+
+    def forward(self, x):
+        out = self.module(x)
+
+        return out
+    

yolo/models/yolov3/yolov3.py

@@ -34,14 +34,14 @@ class Yolov3(nn.Module):
         self.backbone = Yolov3Backbone(cfg)
         self.pyramid_feat_dims = self.backbone.feat_dims[-3:]
         ## Neck: SPP
-        self.neck     = SPPF(cfg, self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1])
+        self.neck = SPPF(self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1])
         self.pyramid_feat_dims[-1] = self.neck.out_dim
         ## Neck: FPN
-        self.fpn      = Yolov3FPN(cfg, self.pyramid_feat_dims)
+        self.fpn = Yolov3FPN(cfg, self.pyramid_feat_dims)
         ## Head
-        self.head     = Yolov3DetHead(cfg, self.fpn.out_dims)
+        self.head = Yolov3DetHead(cfg, self.fpn.out_dims)
         ## Pred
-        self.pred     = Yolov3DetPredLayer(cfg)
+        self.pred = Yolov3DetPredLayer(cfg)
 
     def post_process(self, obj_preds, cls_preds, box_preds):
         """

yolo/models/yolov3/yolov3_backbone.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov3_basic import BasicConv, ResBlock
+    from .modules import ConvModule, ResBlock
 except:
-    from  yolov3_basic import BasicConv, ResBlock
+    from  modules import ConvModule, ResBlock
 
 # IN1K pretrained weight
 pretrained_urls = {
@@ -21,7 +21,7 @@ class Yolov3Backbone(nn.Module):
     def __init__(self, cfg):
         super(Yolov3Backbone, self).__init__()
         # ------------------ Basic setting ------------------
-        self.model_scale = cfg.scale
+        self.model_scale = cfg.model_scale
         self.feat_dims = [round(64   * cfg.width),
                           round(128  * cfg.width),
                           round(256  * cfg.width),
@@ -30,64 +30,46 @@ class Yolov3Backbone(nn.Module):
         
         # ------------------ Network setting ------------------
         ## P1/2
-        self.layer_1 = BasicConv(3, self.feat_dims[0],
-                                 kernel_size=6, padding=2, stride=2,
-                                 act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise)
+        self.layer_1 = ConvModule(3, self.feat_dims[0], kernel_size=6, padding=2, stride=2)
         # P2/4
         self.layer_2 = nn.Sequential(
-            BasicConv(self.feat_dims[0], self.feat_dims[1],
-                      kernel_size=3, padding=1, stride=2,
-                      act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
+            ConvModule(self.feat_dims[0], self.feat_dims[1], kernel_size=3, padding=1, stride=2),
             ResBlock(in_dim     = self.feat_dims[1],
                      out_dim    = self.feat_dims[1],
                      num_blocks = round(3*cfg.depth),
                      expansion  = 0.5,
                      shortcut   = True,
-                     act_type   = cfg.bk_act,
-                     norm_type  = cfg.bk_norm,
-                     depthwise  = cfg.bk_depthwise)
+                     )
         )
         # P3/8
         self.layer_3 = nn.Sequential(
-            BasicConv(self.feat_dims[1], self.feat_dims[2],
-                      kernel_size=3, padding=1, stride=2,
-                      act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
+            ConvModule(self.feat_dims[1], self.feat_dims[2], kernel_size=3, padding=1, stride=2),
             ResBlock(in_dim     = self.feat_dims[2],
                      out_dim    = self.feat_dims[2],
                      num_blocks = round(9*cfg.depth),
                      expansion  = 0.5,
                      shortcut   = True,
-                     act_type   = cfg.bk_act,
-                     norm_type  = cfg.bk_norm,
-                     depthwise  = cfg.bk_depthwise)
+                     )
         )
         # P4/16
         self.layer_4 = nn.Sequential(
-            BasicConv(self.feat_dims[2], self.feat_dims[3],
-                      kernel_size=3, padding=1, stride=2,
-                      act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
+            ConvModule(self.feat_dims[2], self.feat_dims[3], kernel_size=3, padding=1, stride=2),
             ResBlock(in_dim     = self.feat_dims[3],
                      out_dim    = self.feat_dims[3],
                      num_blocks = round(9*cfg.depth),
                      expansion  = 0.5,
                      shortcut   = True,
-                     act_type   = cfg.bk_act,
-                     norm_type  = cfg.bk_norm,
-                     depthwise  = cfg.bk_depthwise)
+                     )
         )
         # P5/32
         self.layer_5 = nn.Sequential(
-            BasicConv(self.feat_dims[3], self.feat_dims[4],
-                      kernel_size=3, padding=1, stride=2,
-                      act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
+            ConvModule(self.feat_dims[3], self.feat_dims[4], kernel_size=3, padding=1, stride=2),
             ResBlock(in_dim     = self.feat_dims[4],
                      out_dim    = self.feat_dims[4],
                      num_blocks = round(3*cfg.depth),
                      expansion  = 0.5,
                      shortcut   = True,
-                     act_type   = cfg.bk_act,
-                     norm_type  = cfg.bk_norm,
-                     depthwise  = cfg.bk_depthwise)
+                     )
         )
 
         # Initialize all layers
@@ -101,8 +83,6 @@ class Yolov3Backbone(nn.Module):
         """Initialize the parameters."""
         for m in self.modules():
             if isinstance(m, torch.nn.Conv2d):
-                # In order to be consistent with the source code,
-                # reset the Conv2d initialization parameters
                 m.reset_parameters()
 
     def load_pretrained(self):
@@ -146,12 +126,9 @@ if __name__ == '__main__':
     from thop import profile
     class BaseConfig(object):
         def __init__(self) -> None:
-            self.bk_act = 'silu'
-            self.bk_norm = 'BN'
-            self.bk_depthwise = False
             self.width = 0.5
             self.depth = 0.34
-            self.scale = "s"
+            self.model_scale = "s"
             self.use_pretrained = True
 
     cfg = BaseConfig()

yolo/models/yolov3/yolov3_basic.py

@@ -1,145 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import List
-
-
-# --------------------- Basic modules ---------------------
-def get_conv2d(c1, c2, k, p, s, d, g, bias=False):
-    conv = nn.Conv2d(c1, c2, k, stride=s, padding=p, dilation=d, groups=g, bias=bias)
-
-    return conv
-
-def get_activation(act_type=None):
-    if act_type == 'relu':
-        return nn.ReLU(inplace=True)
-    elif act_type == 'lrelu':
-        return nn.LeakyReLU(0.1, inplace=True)
-    elif act_type == 'mish':
-        return nn.Mish(inplace=True)
-    elif act_type == 'silu':
-        return nn.SiLU(inplace=True)
-    elif act_type is None:
-        return nn.Identity()
-    else:
-        raise NotImplementedError
-
-def get_norm(norm_type, dim):
-    if norm_type == 'BN':
-        return nn.BatchNorm2d(dim)
-    elif norm_type == 'GN':
-        return nn.GroupNorm(num_groups=32, num_channels=dim)
-    elif norm_type is None:
-        return nn.Identity()
-    else:
-        raise NotImplementedError
-
-class BasicConv(nn.Module):
-    def __init__(self, 
-                 in_dim,                   # in channels
-                 out_dim,                  # out channels 
-                 kernel_size=1,            # kernel size 
-                 padding=0,                # padding
-                 stride=1,                 # padding
-                 dilation=1,               # dilation
-                 act_type  :str = 'lrelu', # activation
-                 norm_type :str = 'BN',    # normalization
-                 depthwise :bool = False
-                ):
-        super(BasicConv, self).__init__()
-        self.depthwise = depthwise
-        use_bias = False if norm_type is not None else True
-        if not depthwise:
-            self.conv = get_conv2d(in_dim, out_dim, k=kernel_size, p=padding, s=stride, d=dilation, g=1, bias=use_bias)
-            self.norm = get_norm(norm_type, out_dim)
-        else:
-            self.conv1 = get_conv2d(in_dim, in_dim, k=kernel_size, p=padding, s=stride, d=dilation, g=in_dim, bias=use_bias)
-            self.norm1 = get_norm(norm_type, in_dim)
-            self.conv2 = get_conv2d(in_dim, out_dim, k=1, p=0, s=1, d=1, g=1)
-            self.norm2 = get_norm(norm_type, out_dim)
-        self.act  = get_activation(act_type)
-
-    def forward(self, x):
-        if not self.depthwise:
-            return self.act(self.norm(self.conv(x)))
-        else:
-            # Depthwise conv
-            x = self.norm1(self.conv1(x))
-            # Pointwise conv
-            x = self.act(self.norm2(self.conv2(x)))
-            return x
-
-
-# ---------------------------- 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
-
-    def forward(self, x):
-        h = self.conv_layer2(self.conv_layer1(x))
-
-        return x + h if self.shortcut else h
-
-class ResBlock(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(ResBlock, self).__init__()
-        # ---------- Basic parameters ----------
-        self.num_blocks = num_blocks
-        self.expansion = expansion
-        self.shortcut = shortcut
-        # ---------- Model parameters ----------
-        module = []
-        for i in range(num_blocks):
-            if i == 0:
-                module.append(YoloBottleneck(in_dim       = in_dim,
-                                             out_dim      = out_dim,
-                                             kernel_size  = [1, 3],
-                                             expansion    = expansion,
-                                             shortcut     = shortcut,
-                                             act_type     = act_type,
-                                             norm_type    = norm_type,
-                                             depthwise    = depthwise))
-            else:
-                module.append(YoloBottleneck(in_dim       = out_dim,
-                                             out_dim      = out_dim,
-                                             kernel_size  = [1, 3],
-                                             expansion    = expansion,
-                                             shortcut     = shortcut,
-                                             act_type     = act_type,
-                                             norm_type    = norm_type,
-                                             depthwise    = depthwise))
-
-
-        self.module = nn.Sequential(*module)
-
-    def forward(self, x):
-        out = self.module(x)
-
-        return out
-    

yolo/models/yolov3/yolov3_fpn.py

@@ -4,9 +4,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 try:
-    from .yolov3_basic import BasicConv, ResBlock
+    from .modules import ConvModule, ResBlock
 except:
-    from  yolov3_basic import BasicConv, ResBlock
+    from  modules import ConvModule, ResBlock
 
 
 # Yolov3FPN
@@ -19,30 +19,22 @@ class Yolov3FPN(nn.Module):
 
         # ---------------------- Yolov3's Top down FPN ----------------------
         ## P5 -> P4
+        self.reduce_layer_1   = ConvModule(round(512*cfg.width), round(256*cfg.width), kernel_size=1, padding=0, stride=1)
         self.top_down_layer_1 = ResBlock(in_dim     = c5,
                                          out_dim    = round(512*cfg.width),
                                          num_blocks = round(3*cfg.depth),
                                          expansion  = 0.5,
                                          shortcut   = False,
-                                         act_type   = cfg.fpn_act,
-                                         norm_type  = cfg.fpn_norm,
-                                         depthwise  = cfg.fpn_depthwise)
-        self.reduce_layer_1   = BasicConv(round(512*cfg.width), round(256*cfg.width),
-                                          kernel_size=1, padding=0, stride=1,
-                                          act_type=cfg.fpn_act, norm_type=cfg.fpn_norm)
+                                         )
 
         ## P4 -> P3
+        self.reduce_layer_2   = ConvModule(round(256*cfg.width), round(128*cfg.width), kernel_size=1, padding=0, stride=1)
         self.top_down_layer_2 = ResBlock(in_dim     = c4 + round(256*cfg.width),
                                          out_dim    = round(256*cfg.width),
                                          num_blocks = round(3*cfg.depth),
                                          expansion  = 0.5,
                                          shortcut   = False,
-                                         act_type   = cfg.fpn_act,
-                                         norm_type  = cfg.fpn_norm,
-                                         depthwise  = cfg.fpn_depthwise)
-        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)
+                                         )
         
         ## P3
         self.top_down_layer_3 = ResBlock(in_dim     = c3 + round(128*cfg.width),
@@ -50,14 +42,11 @@ class Yolov3FPN(nn.Module):
                                          num_blocks = round(3*cfg.depth),
                                          expansion  = 0.5,
                                          shortcut   = False,
-                                         act_type   = cfg.fpn_act,
-                                         norm_type  = cfg.fpn_norm,
-                                         depthwise  = cfg.fpn_depthwise)
+                                         )
 
         # ---------------------- Yolov3's output projection ----------------------
         self.out_layers = nn.ModuleList([
-            BasicConv(in_dim, round(cfg.head_dim*cfg.width), kernel_size=1,
-                      act_type=cfg.fpn_act, norm_type=cfg.fpn_norm)
+            ConvModule(in_dim, round(cfg.head_dim*cfg.width), kernel_size=1)
                       for in_dim in [round(128*cfg.width), round(256*cfg.width), round(512*cfg.width)]
                       ])
         self.out_dims = [round(cfg.head_dim*cfg.width)] * 3
@@ -69,8 +58,6 @@ class Yolov3FPN(nn.Module):
         """Initialize the parameters."""
         for m in self.modules():
             if isinstance(m, torch.nn.Conv2d):
-                # In order to be consistent with the source code,
-                # reset the Conv2d initialization parameters
                 m.reset_parameters()
 
     def forward(self, features):
@@ -111,10 +98,6 @@ if __name__=='__main__':
             self.out_stride = [8, 16, 32]
             self.max_stride = 32
             self.num_levels = 3
-            ## FPN
-            self.fpn_act  = 'silu'
-            self.fpn_norm = 'BN'
-            self.fpn_depthwise = False
             ## Head
             self.head_dim = 256
 

yolo/models/yolov3/yolov3_head.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov3_basic import BasicConv
+    from .modules import ConvModule
 except:
-    from  yolov3_basic import BasicConv
+    from  modules import ConvModule
 
 
 ## Single-level Detection Head
@@ -15,17 +15,12 @@ class DetHead(nn.Module):
                  reg_head_dim :int  = 256,
                  num_cls_head :int  = 2,
                  num_reg_head :int  = 2,
-                 act_type     :str  = "silu",
-                 norm_type    :str  = "BN",
-                 depthwise    :bool = False):
+                 ):
         super().__init__()
         # --------- Basic Parameters ----------
         self.in_dim = in_dim
         self.num_cls_head = num_cls_head
         self.num_reg_head = num_reg_head
-        self.act_type = act_type
-        self.norm_type = norm_type
-        self.depthwise = depthwise
         
         # --------- Network Parameters ----------
         ## cls head
@@ -33,54 +28,20 @@ class DetHead(nn.Module):
         self.cls_head_dim = cls_head_dim
         for i in range(num_cls_head):
             if i == 0:
-                cls_feats.append(
-                    BasicConv(in_dim, self.cls_head_dim,
-                              kernel_size=3, padding=1, stride=1, 
-                              act_type=act_type,
-                              norm_type=norm_type,
-                              depthwise=depthwise)
-                              )
+                cls_feats.append(ConvModule(in_dim, self.cls_head_dim, kernel_size=3, padding=1, stride=1))
             else:
-                cls_feats.append(
-                    BasicConv(self.cls_head_dim, self.cls_head_dim,
-                              kernel_size=3, padding=1, stride=1, 
-                              act_type=act_type,
-                              norm_type=norm_type,
-                              depthwise=depthwise)
-                              )
+                cls_feats.append(ConvModule(self.cls_head_dim, self.cls_head_dim, kernel_size=3, padding=1, stride=1))
         ## reg head
         reg_feats = []
         self.reg_head_dim = reg_head_dim
         for i in range(num_reg_head):
             if i == 0:
-                reg_feats.append(
-                    BasicConv(in_dim, self.reg_head_dim,
-                              kernel_size=3, padding=1, stride=1, 
-                              act_type=act_type,
-                              norm_type=norm_type,
-                              depthwise=depthwise)
-                              )
+                reg_feats.append(ConvModule(in_dim, self.reg_head_dim, kernel_size=3, padding=1, stride=1))
             else:
-                reg_feats.append(
-                    BasicConv(self.reg_head_dim, self.reg_head_dim,
-                              kernel_size=3, padding=1, stride=1, 
-                              act_type=act_type,
-                              norm_type=norm_type,
-                              depthwise=depthwise)
-                              )
+                reg_feats.append(ConvModule(self.reg_head_dim, self.reg_head_dim, kernel_size=3, padding=1, stride=1))
         self.cls_feats = nn.Sequential(*cls_feats)
         self.reg_feats = nn.Sequential(*reg_feats)
 
-        self.init_weights()
-        
-    def init_weights(self):
-        """Initialize the parameters."""
-        for m in self.modules():
-            if isinstance(m, torch.nn.Conv2d):
-                # In order to be consistent with the source code,
-                # reset the Conv2d initialization parameters
-                m.reset_parameters()
-
     def forward(self, x):
         """
             in_feats: (Tensor) [B, C, H, W]
@@ -101,16 +62,23 @@ class Yolov3DetHead(nn.Module):
                      reg_head_dim = round(cfg.head_dim * cfg.width),
                      num_cls_head = cfg.num_cls_head,
                      num_reg_head = cfg.num_reg_head,
-                     act_type     = cfg.head_act,
-                     norm_type    = cfg.head_norm,
-                     depthwise    = cfg.head_depthwise)
-                     for level in range(cfg.num_levels)
-                     ])
+                     ) for level in range(len(cfg.out_stride))])
         # --------- Basic Parameters ----------
         self.in_dims = in_dims
         self.cls_head_dim = cfg.head_dim
         self.reg_head_dim = cfg.head_dim
 
+        # Initialize all layers
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialize the parameters."""
+        for m in self.modules():
+            if isinstance(m, torch.nn.Conv2d):
+                # In order to be consistent with the source code,
+                # reset the Conv2d initialization parameters
+                m.reset_parameters()
+
     def forward(self, feats):
         """
             feats: List[(Tensor)] [[B, C, H, W], ...]
@@ -142,9 +110,6 @@ if __name__=='__main__':
             self.max_stride = 32
             self.num_levels = 3
             ## Head
-            self.head_act  = 'lrelu'
-            self.head_norm = 'BN'
-            self.head_depthwise = False
             self.head_dim  = 256
             self.num_cls_head   = 2
             self.num_reg_head   = 2
@@ -169,4 +134,4 @@ if __name__=='__main__':
     flops, params = profile(head, inputs=(pyramid_feats, ), verbose=False)
     print('==============================')
     print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Params : {:.2f} M'.format(params / 1e6))    
+    print('Params : {:.2f} M'.format(params / 1e6))

yolo/models/yolov3/yolov3_neck.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov3_basic import BasicConv
+    from .modules import ConvModule
 except:
-    from  yolov3_basic import BasicConv
+    from  modules import ConvModule
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
@@ -12,21 +12,15 @@ class SPPF(nn.Module):
     """
         This code referenced to https://github.com/ultralytics/yolov5
     """
-    def __init__(self, cfg, in_dim, out_dim):
+    def __init__(self, in_dim, out_dim):
         super().__init__()
         ## ----------- Basic Parameters -----------
-        inter_dim = round(in_dim * cfg.neck_expand_ratio)
+        inter_dim = in_dim // 2
         self.out_dim = out_dim
         ## ----------- Network Parameters -----------
-        self.cv1 = BasicConv(in_dim, inter_dim,
-                             kernel_size=1, padding=0, stride=1,
-                             act_type=cfg.neck_act, norm_type=cfg.neck_norm)
-        self.cv2 = BasicConv(inter_dim * 4, out_dim,
-                             kernel_size=1, padding=0, stride=1,
-                             act_type=cfg.neck_act, norm_type=cfg.neck_norm)
-        self.m = nn.MaxPool2d(kernel_size=cfg.spp_pooling_size,
-                              stride=1,
-                              padding=cfg.spp_pooling_size // 2)
+        self.cv1 = ConvModule(in_dim, inter_dim, kernel_size=1, padding=0, stride=1)
+        self.cv2 = ConvModule(inter_dim * 4, out_dim, kernel_size=1, padding=0, stride=1)
+        self.m = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)
 
         # Initialize all layers
         self.init_weights()
@@ -35,8 +29,6 @@ class SPPF(nn.Module):
         """Initialize the parameters."""
         for m in self.modules():
             if isinstance(m, torch.nn.Conv2d):
-                # In order to be consistent with the source code,
-                # reset the Conv2d initialization parameters
                 m.reset_parameters()
 
     def forward(self, x):
@@ -52,24 +44,10 @@ if __name__=='__main__':
     from thop import profile
     # Model config
     
-    # YOLOv3-Base config
-    class Yolov3BaseConfig(object):
-        def __init__(self) -> None:
-            # ---------------- Model config ----------------
-            self.out_stride = 32
-            self.max_stride = 32
-            ## Neck
-            self.neck_act       = 'lrelu'
-            self.neck_norm      = 'BN'
-            self.neck_depthwise = False
-            self.neck_expand_ratio = 0.5
-            self.spp_pooling_size  = 5
-
-    cfg = Yolov3BaseConfig()
-    # Build a head
+    # Build a neck
     in_dim  = 512
     out_dim = 512
-    neck = SPPF(cfg, in_dim, out_dim)
+    neck = SPPF(in_dim, out_dim)
 
     # Inference
     x = torch.randn(1, in_dim, 20, 20)

yolo/models/yolov3/yolov3_pred.py

@@ -116,6 +116,7 @@ class Yolov3DetPredLayer(nn.Module):
         super().__init__()
         # --------- Basic Parameters ----------
         self.cfg = cfg
+        self.num_levels = len(cfg.out_stride)
 
         # ----------- Network Parameters -----------
         ## pred layers
@@ -125,7 +126,7 @@ class Yolov3DetPredLayer(nn.Module):
                           stride       = cfg.out_stride[level],
                           anchor_sizes = cfg.anchor_size[level],
                           num_classes  = cfg.num_classes,)
-                          for level in range(cfg.num_levels)
+                          for level in range(self.num_levels)
                           ])
 
     def forward(self, cls_feats, reg_feats):
@@ -136,7 +137,7 @@ class Yolov3DetPredLayer(nn.Module):
         all_cls_preds = []
         all_reg_preds = []
         all_box_preds = []
-        for level in range(self.cfg.num_levels):
+        for level in range(self.num_levels):
             # -------------- Single-level prediction --------------
             outputs = self.multi_level_preds[level](cls_feats[level], reg_feats[level])