modify v5

yolo/models/yolov5/modules.py

@@ -0,0 +1,78 @@
+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)))
+
+
+# ---------------------------- 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,
+                 ) -> 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 CSPBlock(nn.Module):
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 num_blocks :int   = 1,
+                 expansion  :float = 0.5,
+                 shortcut   :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 = ConvModule(in_dim, inter_dim, kernel_size=1)
+        self.conv_layer_2 = ConvModule(in_dim, inter_dim, kernel_size=1)
+        self.conv_layer_3 = ConvModule(inter_dim * 2, out_dim, kernel_size=1)
+        self.module = nn.Sequential(*[
+            YoloBottleneck(inter_dim,
+                           inter_dim,
+                           kernel_size = [1, 3],
+                           expansion   = 1.0,
+                           shortcut    = shortcut,
+                           ) for _ in range(num_blocks)])
+
+    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))
+
+        return out
+    

yolo/models/yolov5/yolov5.py

@@ -34,14 +34,14 @@ class Yolov5(nn.Module):
         self.backbone = Yolov5Backbone(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      = Yolov5PaFPN(cfg, self.pyramid_feat_dims)
+        self.fpn = Yolov5PaFPN(cfg, self.pyramid_feat_dims)
         ## Head
-        self.head     = Yolov5DetHead(cfg, self.fpn.out_dims)
+        self.head = Yolov5DetHead(cfg, self.fpn.out_dims)
         ## Pred
-        self.pred     = Yolov5DetPredLayer(cfg)
+        self.pred = Yolov5DetPredLayer(cfg)
 
     def post_process(self, obj_preds, cls_preds, box_preds):
         """

yolo/models/yolov5/yolov5_backbone.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov5_basic import BasicConv, CSPBlock
+    from .modules import ConvModule, CSPBlock
 except:
-    from  yolov5_basic import BasicConv, CSPBlock
+    from  modules import ConvModule, CSPBlock
 
 # IN1K pretrained weight
 pretrained_urls = {
@@ -31,64 +31,46 @@ class Yolov5Backbone(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),
             CSPBlock(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),
             CSPBlock(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),
             CSPBlock(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),
             CSPBlock(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
@@ -102,8 +84,6 @@ class Yolov5Backbone(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):
@@ -147,9 +127,6 @@ 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"

yolo/models/yolov5/yolov5_basic.py

@@ -1,138 +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 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)
-                                                           ])
-
-    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))
-
-        return out
-    

yolo/models/yolov5/yolov5_head.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov5_basic import BasicConv
+    from .modules import ConvModule
 except:
-    from  yolov5_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,41 +28,17 @@ 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)
 
@@ -91,11 +62,7 @@ class Yolov5DetHead(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(cfg.num_levels)])
         # --------- Basic Parameters ----------
         self.in_dims = in_dims
         self.cls_head_dim = cfg.head_dim
@@ -143,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

yolo/models/yolov5/yolov5_neck.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov5_basic import BasicConv
+    from .modules import ConvModule
 except:
-    from  yolov5_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
     
-    # YOLOv5-Base config
-    class Yolov5BaseConfig(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 = Yolov5BaseConfig()
-    # 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/yolov5/yolov5_pafpn.py

@@ -4,9 +4,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 try:
-    from .yolov5_basic import BasicConv, CSPBlock
+    from .modules import ConvModule, CSPBlock
 except:
-    from  yolov5_basic import BasicConv, CSPBlock
+    from  modules import ConvModule, CSPBlock
 
 
 # Yolov5FPN
@@ -18,61 +18,44 @@ class Yolov5PaFPN(nn.Module):
 
         # ---------------------- Yolov5's Top down FPN ----------------------
         ## P5 -> P4
-        self.reduce_layer_1   = BasicConv(c5, round(512*cfg.width),
-                                          kernel_size=1, padding=0, stride=1,
-                                          act_type=cfg.fpn_act, norm_type=cfg.fpn_norm)
+        self.reduce_layer_1   = ConvModule(c5, round(512*cfg.width), kernel_size=1, padding=0, stride=1)
         self.top_down_layer_1 = CSPBlock(in_dim     = c4 + round(512*cfg.width),
                                          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)
+                                         )
 
         ## P4 -> P3
-        self.reduce_layer_2   = 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)
+        self.reduce_layer_2   = ConvModule(round(512*cfg.width), round(256*cfg.width), kernel_size=1, padding=0, stride=1)
         self.top_down_layer_2 = CSPBlock(in_dim     = c3 + 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)
+                                         )
         
         # ---------------------- Yolov5's Bottom up PAN ----------------------
         ## P3 -> P4
-        self.downsample_layer_1 = 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.downsample_layer_1 = ConvModule(round(256*cfg.width), round(256*cfg.width), kernel_size=3, padding=1, stride=2)
         self.bottom_up_layer_1  = CSPBlock(in_dim     = round(256*cfg.width) + round(256*cfg.width),
                                            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)
+                                           )
         ## P4 -> P5
-        self.downsample_layer_2 = BasicConv(round(512*cfg.width), round(512*cfg.width),
-                                            kernel_size=3, padding=1, stride=2,
-                                            act_type=cfg.fpn_act, norm_type=cfg.fpn_norm, depthwise=cfg.fpn_depthwise)
+        self.downsample_layer_2 = ConvModule(round(512*cfg.width), round(512*cfg.width), kernel_size=3, padding=1, stride=2)
         self.bottom_up_layer_2  = CSPBlock(in_dim     = round(512*cfg.width) + round(512*cfg.width),
                                            out_dim    = round(1024*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)
+                                           )
 
         # ---------------------- Yolov5'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(256*cfg.width), round(512*cfg.width), round(1024*cfg.width)]
                       ])
         self.out_dims = [round(cfg.head_dim*cfg.width)] * 3
@@ -84,8 +67,6 @@ class Yolov5PaFPN(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):
@@ -135,10 +116,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