Browse Source

add YOLOv3-Tiny

yjh0410 2 năm trước cách đây
mục cha
commit
38baee3b6e

+ 2 - 2
config/__init__.py

@@ -19,8 +19,8 @@ def build_model_config(args):
     elif args.model == 'yolov2':
         cfg = yolov2_cfg
     # YOLOv3
-    elif args.model == 'yolov3':
-        cfg = yolov3_cfg
+    elif args.model in ['yolov3', 'yolov3_t']:
+        cfg = yolov3_cfg[args.model]
     # YOLOv4
     elif args.model == 'yolov4':
         cfg = yolov4_cfg

+ 122 - 53
config/yolov3_config.py

@@ -1,57 +1,126 @@
 # YOLOv3 Config
 
 yolov3_cfg = {
-    # input
-    'trans_type': 'yolov5_large',
-    'multi_scale': [0.5, 1.0],
-    # model
-    'backbone': 'darknet53',
-    'pretrained': True,
-    'stride': [8, 16, 32],  # P3, P4, P5
-    'width': 1.0,
-    'depth': 1.0,
-    # neck
-    'neck': 'sppf',
-    'expand_ratio': 0.5,
-    'pooling_size': 5,
-    'neck_act': 'silu',
-    'neck_norm': 'BN',
-    'neck_depthwise': False,
-     # fpn
-    'fpn': 'yolo_fpn',
-    'fpn_act': 'silu',
-    'fpn_norm': 'BN',
-    'fpn_depthwise': False,
-    # head
-    'head': 'decoupled_head',
-    'head_act': 'silu',
-    'head_norm': 'BN',
-    'num_cls_head': 2,
-    'num_reg_head': 2,
-    'head_depthwise': False,
-    'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
-                    [30, 61],   [62, 45],   [59, 119],    # P4
-                    [116, 90],  [156, 198], [373, 326]],  # P5
-    # matcher
-    'iou_thresh': 0.5,
-    # loss weight
-    'loss_obj_weight': 1.0,
-    'loss_cls_weight': 1.0,
-    'loss_box_weight': 5.0,
-    # training configuration
-    'no_aug_epoch': 10,
-    # optimizer
-    'optimizer': 'sgd',        # optional: sgd, adam, adamw
-    'momentum': 0.937,         # SGD: 0.937;    AdamW: invalid
-    'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
-    'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
-    # model EMA
-    'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
-    'ema_tau': 2000,
-    # lr schedule
-    'scheduler': 'linear',
-    'lr0': 0.01,               # SGD: 0.01;     AdamW: 0.004
-    'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.05
-    'warmup_momentum': 0.8,
-    'warmup_bias_lr': 0.1,
+    'yolov3':{
+        # ---------------- Model config ----------------
+        ## Backbone
+        'backbone': 'darknet53',
+        'pretrained': True,
+        'stride': [8, 16, 32],  # P3, P4, P5
+        'width': 1.0,
+        'depth': 1.0,
+        ## Neck
+        'neck': 'sppf',
+        'expand_ratio': 0.5,
+        'pooling_size': 5,
+        'neck_act': 'silu',
+        'neck_norm': 'BN',
+        'neck_depthwise': False,
+        ## FPN
+        'fpn': 'yolo_fpn',
+        'fpn_act': 'silu',
+        'fpn_norm': 'BN',
+        'fpn_depthwise': False,
+        ## Head
+        'head': 'decoupled_head',
+        'head_act': 'silu',
+        'head_norm': 'BN',
+        'num_cls_head': 2,
+        'num_reg_head': 2,
+        'head_depthwise': False,
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        # ---------------- Train config ----------------
+        ## input
+        'trans_type': 'yolov5_large',
+        'multi_scale': [0.5, 1.0],
+        # ---------------- Assignment config ----------------
+        ## matcher
+        'iou_thresh': 0.5,
+        # ---------------- Loss config ----------------
+        ## loss weight
+        'loss_obj_weight': 1.0,
+        'loss_cls_weight': 1.0,
+        'loss_box_weight': 5.0,
+        # ---------------- Train config ----------------
+        ## close strong augmentation
+        'no_aug_epoch': 10,
+        ## optimizer
+        'optimizer': 'sgd',        # optional: sgd, AdamW
+        'momentum': 0.937,         # SGD: 0.937;    AdamW: None
+        'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
+        'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
+        ## model EMA
+        'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
+        'ema_tau': 2000,
+        ## lr schedule
+        'scheduler': 'linear',
+        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.001
+        'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.01
+        'warmup_momentum': 0.8,
+        'warmup_bias_lr': 0.1,
+    },
+
+    'yolov3_t':{
+        # ---------------- Model config ----------------
+        ## Backbone
+        'backbone': 'darknet_tiny',
+        'pretrained': True,
+        'stride': [8, 16, 32],  # P3, P4, P5
+        'width': 0.25,
+        'depth': 0.34,
+        ## Neck
+        'neck': 'sppf',
+        'expand_ratio': 0.5,
+        'pooling_size': 5,
+        'neck_act': 'silu',
+        'neck_norm': 'BN',
+        'neck_depthwise': False,
+        ## FPN
+        'fpn': 'yolo_fpn',
+        'fpn_act': 'silu',
+        'fpn_norm': 'BN',
+        'fpn_depthwise': False,
+        ## Head
+        'head': 'decoupled_head',
+        'head_act': 'silu',
+        'head_norm': 'BN',
+        'num_cls_head': 2,
+        'num_reg_head': 2,
+        'head_depthwise': False,
+        'anchor_size': [[10, 13],   [16, 30],   [33, 23],     # P3
+                        [30, 61],   [62, 45],   [59, 119],    # P4
+                        [116, 90],  [156, 198], [373, 326]],  # P5
+        # ---------------- Train config ----------------
+        ## input
+        'trans_type': 'yolov5_tiny',
+        'multi_scale': [0.5, 1.0],
+        # ---------------- Assignment config ----------------
+        ## matcher
+        'iou_thresh': 0.5,
+        # ---------------- Loss config ----------------
+        ## loss weight
+        'loss_obj_weight': 1.0,
+        'loss_cls_weight': 1.0,
+        'loss_box_weight': 5.0,
+        # ---------------- Train config ----------------
+        ## close strong augmentation
+        'no_aug_epoch': 10,
+        ## optimizer
+        'optimizer': 'sgd',        # optional: sgd, AdamW
+        'momentum': 0.937,         # SGD: 0.937;    AdamW: None
+        'weight_decay': 5e-4,      # SGD: 5e-4;     AdamW: 5e-2
+        'clip_grad': 10,           # SGD: 10.0;     AdamW: -1
+        ## model EMA
+        'ema_decay': 0.9999,       # SGD: 0.9999;   AdamW: 0.9998
+        'ema_tau': 2000,
+        ## lr schedule
+        'scheduler': 'linear',
+        'lr0': 0.01,              # SGD: 0.01;     AdamW: 0.001
+        'lrf': 0.01,               # SGD: 0.01;     AdamW: 0.01
+        'warmup_momentum': 0.8,
+        'warmup_bias_lr': 0.1,
+    },
+
 }

+ 1 - 1
models/detectors/__init__.py

@@ -27,7 +27,7 @@ def build_model(args,
         model, criterion = build_yolov2(
             args, model_cfg, device, num_classes, trainable)
     # YOLOv3   
-    elif args.model == 'yolov3':
+    elif args.model in ['yolov3', 'yolov3_t']:
         model, criterion = build_yolov3(
             args, model_cfg, device, num_classes, trainable)
     # YOLOv4   

+ 50 - 1
models/detectors/yolov3/yolov3_backbone.py

@@ -8,11 +8,13 @@ except:
     
 
 model_urls = {
+    "darknet_tiny": None,
     "darknet53": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet53_silu.pth",
 }
 
 
 # --------------------- DarkNet-53 -----------------------
+## DarkNet-53
 class DarkNet53(nn.Module):
     def __init__(self, act_type='silu', norm_type='BN'):
         super(DarkNet53, self).__init__()
@@ -57,6 +59,50 @@ class DarkNet53(nn.Module):
 
         return outputs
 
+## DarkNet-Tiny
+class DarkNetTiny(nn.Module):
+    def __init__(self, act_type='silu', norm_type='BN'):
+        super(DarkNetTiny, self).__init__()
+        self.feat_dims = [64, 128, 256]
+
+        # stride = 2
+        self.layer_1 = nn.Sequential(
+            Conv(3, 16, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
+            ResBlock(16, 16, nblocks=1, act_type=act_type, norm_type=norm_type)
+        )
+        # stride = 4
+        self.layer_2 = nn.Sequential(
+            Conv(16, 32, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
+            ResBlock(32, 32, nblocks=1, act_type=act_type, norm_type=norm_type)
+        )
+        # stride = 8
+        self.layer_3 = nn.Sequential(
+            Conv(32, 64, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
+            ResBlock(64, 64, nblocks=3, act_type=act_type, norm_type=norm_type)
+        )
+        # stride = 16
+        self.layer_4 = nn.Sequential(
+            Conv(64, 128, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
+            ResBlock(128, 128, nblocks=3, act_type=act_type, norm_type=norm_type)
+        )
+        # stride = 32
+        self.layer_5 = nn.Sequential(
+            Conv(128, 256, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
+            ResBlock(256, 256, nblocks=2, act_type=act_type, norm_type=norm_type)
+        )
+
+
+    def forward(self, x):
+        c1 = self.layer_1(x)
+        c2 = self.layer_2(c1)
+        c3 = self.layer_3(c2)
+        c4 = self.layer_4(c3)
+        c5 = self.layer_5(c4)
+
+        outputs = [c3, c4, c5]
+
+        return outputs
+
 
 # --------------------- Functions -----------------------
 def build_backbone(model_name='darknet53', pretrained=False): 
@@ -67,9 +113,12 @@ def build_backbone(model_name='darknet53', pretrained=False):
     if model_name == 'darknet53':
         backbone = DarkNet53(act_type='silu', norm_type='BN')
         feat_dims = backbone.feat_dims
+    elif model_name == 'darknet_tiny':
+        backbone = DarkNetTiny(act_type='silu', norm_type='BN')
+        feat_dims = backbone.feat_dims
 
     if pretrained:
-        url = model_urls['darknet53']
+        url = model_urls[model_name]
         if url is not None:
             print('Loading pretrained weight ...')
             checkpoint = torch.hub.load_state_dict_from_url(