modify yolov1

yolo/config/yolov1_config.py

@@ -4,6 +4,8 @@
 def build_yolov1_config(args):
     if args.model == 'yolov1_r18':
         return Yolov1R18Config()
+    if args.model == 'yolov1_r50':
+        return Yolov1R50Config()
     else:
         raise NotImplementedError("No config for model: {}".format(args.model))
     
@@ -14,21 +16,12 @@ class Yolov1BaseConfig(object):
         self.out_stride = 32
         self.max_stride = 32
         ## Backbone
-        self.backbone       = 'resnet50'
+        self.backbone = 'resnet50'
         self.use_pretrained = True
-        ## Neck
-        self.neck_act       = 'lrelu'
-        self.neck_norm      = 'BN'
-        self.neck_depthwise = False
-        self.neck_expand_ratio = 0.5
-        self.spp_pooling_size  = 5
         ## Head
-        self.head_act  = 'lrelu'
-        self.head_norm = 'BN'
-        self.head_depthwise = False
         self.head_dim  = 512
-        self.num_cls_head   = 2
-        self.num_reg_head   = 2
+        self.num_cls_head = 2
+        self.num_reg_head = 2
 
         # ---------------- Post-process config ----------------
         ## Post process

yolo/models/yolov1/yolov1_basic.py → yolo/models/yolov1/modules.py

@@ -4,35 +4,6 @@ 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
@@ -41,32 +12,14 @@ class BasicConv(nn.Module):
                  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)
+        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.LeakyReLU(0.1, inplace=True)
 
     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
+        return self.act(self.norm(self.conv(x)))
 
 
 # --------------------- ResNet modules ---------------------

yolo/models/yolov1/resnet.py

@@ -3,9 +3,9 @@ import torch.nn as nn
 import torch.utils.model_zoo as model_zoo
 
 try:
-    from .yolov1_basic import conv1x1, BasicBlock, Bottleneck
+    from .modules import conv1x1, BasicBlock, Bottleneck
 except:
-    from  yolov1_basic import conv1x1, BasicBlock, Bottleneck
+    from  modules import conv1x1, BasicBlock, Bottleneck
 
 __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
            'resnet152']

yolo/models/yolov1/yolov1.py

@@ -30,7 +30,7 @@ class Yolov1(nn.Module):
         
         # ---------------------- Network Parameters ----------------------
         self.backbone = Yolov1Backbone(cfg)
-        self.neck     = SPPF(cfg, self.backbone.feat_dim, cfg.head_dim)
+        self.neck     = SPPF(self.backbone.feat_dim, cfg.head_dim)
         self.head     = Yolov1DetHead(cfg, self.neck.out_dim)
         self.pred     = Yolov1DetPredLayer(cfg)
 

yolo/models/yolov1/yolov1_head.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov1_basic import BasicConv
+    from .modules import BasicConv
 except:
-    from  yolov1_basic import BasicConv
+    from  modules import BasicConv
 
 
 class Yolov1DetHead(nn.Module):
@@ -16,60 +16,30 @@ class Yolov1DetHead(nn.Module):
         self.reg_head_dim = cfg.head_dim
         self.num_cls_head = cfg.num_cls_head
         self.num_reg_head = cfg.num_reg_head
-        self.act_type     = cfg.head_act
-        self.norm_type    = cfg.head_norm
-        self.depthwise    = cfg.head_depthwise
         
         # --------- Network Parameters ----------
         ## cls head
         cls_feats = []
         for i in range(self.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  = self.act_type,
-                              norm_type = self.norm_type,
-                              depthwise = self.depthwise)
-                              )
+                cls_feats.append(BasicConv(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  = self.act_type,
-                              norm_type = self.norm_type,
-                              depthwise = self.depthwise)
-                              )
+                cls_feats.append(BasicConv(self.cls_head_dim, self.cls_head_dim, kernel_size=3, padding=1, stride=1))
         ## reg head
         reg_feats = []
         for i in range(self.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  = self.act_type,
-                              norm_type = self.norm_type,
-                              depthwise = self.depthwise)
-                              )
+                reg_feats.append(BasicConv(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  = self.act_type,
-                              norm_type = self.norm_type,
-                              depthwise = self.depthwise)
-                              )
+                reg_feats.append(BasicConv(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):
@@ -92,9 +62,6 @@ if __name__=='__main__':
             self.out_stride = 32
             self.max_stride = 32
             ## 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/yolov1/yolov1_neck.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolov1_basic import BasicConv
+    from .modules import BasicConv
 except:
-    from  yolov1_basic import BasicConv
+    from  modules import BasicConv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
@@ -12,31 +12,22 @@ 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 = BasicConv(in_dim, inter_dim, kernel_size=1, padding=0, stride=1)
+        self.cv2 = BasicConv(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()
 
     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):
@@ -57,9 +48,6 @@ if __name__=='__main__':
             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 = Yolov1BaseConfig()