use yolov5 style augmentation

config/__init__.py

@@ -123,7 +123,7 @@ def build_model_config(args):
     elif args.model in ['yolox_n_adamw', 'yolox_s_adamw', 'yolox_m_adamw', 'yolox_l_adamw', 'yolox_x_adamw']:
         cfg = yolox_adamw_cfg[args.model]
     # RTCDet
-    elif args.model in ['rtcdet_n', 'rtcdet_s', 'rtcdet_m', 'rtcdet_l', 'rtcdet_x']:
+    elif args.model in ['rtcdet_n', 'rtcdet_t', 'rtcdet_s', 'rtcdet_m', 'rtcdet_l', 'rtcdet_x']:
         cfg = rtcdet_cfg[args.model]
     # CenterNet
     elif args.model in ['ctrnet_n', 'ctrnet_s', 'ctrnet_m', 'ctrnet_l', 'ctrnet_x']:

config/model_config/rtcdet_config.py

@@ -7,7 +7,6 @@ rtcdet_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -58,11 +57,64 @@ rtcdet_cfg = {
         'trainer_type': 'rtcdet',
     },
 
+    'rtcdet_t':{
+        # ---------------- Model config ----------------
+        ## Backbone
+        'bk_pretrained': True,
+        'bk_act': 'silu',
+        'bk_norm': 'BN',
+        'bk_depthwise': False,
+        'width': 0.375,
+        'depth': 0.34,
+        'ratio': 2.0,
+        'stride': [8, 16, 32],  # P3, P4, P5
+        'max_stride': 32,
+        ## Neck: SPP
+        'neck': 'sppf',
+        'neck_expand_ratio': 0.5,
+        'pooling_size': 5,
+        'neck_act': 'silu',
+        'neck_norm': 'BN',
+        'neck_depthwise': False,
+        ## Neck: PaFPN
+        'fpn': 'rtcdet_pafpn',
+        'fpn_act': 'silu',
+        'fpn_norm': 'BN',
+        'fpn_depthwise': False,
+        ## Head
+        'det_head': {'name': 'decoupled_head',
+                     'num_cls_head': 2,
+                     'num_reg_head': 2,
+                     'head_act': 'silu',
+                     'head_norm': 'BN',
+                     'head_depthwise': False,  
+                     },
+        'seg_head': {'name': None,
+                     },
+        'pos_head': {'name': None,
+                     },
+        # ---------------- Train config ----------------
+        ## input
+        'multi_scale': [0.5, 1.25],   # 320 -> 800
+        'trans_type': 'yolov5_n',
+        # ---------------- Assignment config ----------------
+        ## Matcher
+        'matcher': "aligned_simota",
+        'matcher_hpy': {'soft_center_radius': 3.0,
+                        'topk_candidates': 13},
+        # ---------------- Loss config ----------------
+        ## loss weight
+        'loss_cls_weight': 1.0,
+        'loss_box_weight': 2.0,
+        'loss_box_aux': False,
+        # ---------------- Train config ----------------
+        'trainer_type': 'rtcdet',
+    },
+
     'rtcdet_s':{
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -117,7 +169,6 @@ rtcdet_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -172,7 +223,6 @@ rtcdet_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -227,7 +277,6 @@ rtcdet_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -287,7 +336,6 @@ rtcdet_seg_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -347,7 +395,6 @@ rtcdet_pos_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,
@@ -407,7 +454,6 @@ rtcdet_seg_pos_cfg = {
         # ---------------- Model config ----------------
         ## Backbone
         'bk_pretrained': True,
-        'bk_pretrained_mae': False,
         'bk_act': 'silu',
         'bk_norm': 'BN',
         'bk_depthwise': False,

models/detectors/__init__.py

@@ -64,7 +64,7 @@ def build_model(args,
         model, criterion = build_yolox(
             args, model_cfg, device, num_classes, trainable, deploy)
     # RTCDet
-    elif args.model in ['rtcdet_n', 'rtcdet_s', 'rtcdet_m', 'rtcdet_l', 'rtcdet_x']:
+    elif args.model in ['rtcdet_n', 'rtcdet_t', 'rtcdet_s', 'rtcdet_m', 'rtcdet_l', 'rtcdet_x']:
         model, criterion = build_rtcdet(
             args, model_cfg, device, num_classes, trainable, deploy)
     # CenterNet

models/detectors/rtcdet/rtcdet.py

@@ -30,35 +30,37 @@ class RTCDet(nn.Module):
                  nms_class_agnostic = False,
                  ):
         super(RTCDet, self).__init__()
-        # ---------------- Basic Parameters ----------------
+        # ---------------- Basic settings ----------------
+        ## Basic parameters
         self.cfg = cfg
         self.device = device
-        self.strides = cfg['stride']
-        self.num_classes = num_classes
+        self.deploy = deploy
         self.trainable = trainable
-        self.conf_thresh = conf_thresh
-        self.nms_thresh = nms_thresh
-        self.num_levels = len(self.strides)
         self.num_classes = num_classes
+        ## Network parameters
+        self.strides = cfg['stride']
+        self.num_levels = len(self.strides)
+        self.head_dim = round(256 * cfg['width'])
+        ## Post-process parameters
+        self.nms_thresh = nms_thresh
+        self.conf_thresh = conf_thresh
         self.topk_candidates = topk
-        self.deploy = deploy
         self.no_multi_labels = no_multi_labels
         self.nms_class_agnostic = nms_class_agnostic
-        self.head_dim = round(256 * cfg['width'])
         
-        # ---------------- Network Parameters ----------------
-        ## ----------- Backbone -----------
-        self.backbone, feat_dims = build_backbone(cfg, pretrained=cfg['bk_pretrained']&trainable)
+        # ---------------- Network settings ----------------
+        ## Backbone
+        self.backbone, self.fpn_feat_dims = build_backbone(cfg, pretrained=cfg['bk_pretrained']&trainable)
 
-        ## ----------- Neck: SPP -----------
-        self.neck = build_neck(cfg, feat_dims[-1], feat_dims[-1])
-        feat_dims[-1] = self.neck.out_dim
+        ## Neck: SPP
+        self.neck = build_neck(cfg, self.fpn_feat_dims[-1], self.fpn_feat_dims[-1])
+        self.fpn_feat_dims[-1] = self.neck.out_dim
         
-        ## ----------- Neck: FPN -----------
-        self.fpn = build_fpn(cfg, feat_dims, out_dim=self.head_dim)
+        ## Neck: FPN
+        self.fpn = build_fpn(cfg, self.fpn_feat_dims, self.head_dim)
         self.fpn_dims = self.fpn.out_dim
 
-        ## ----------- Head -----------
+        ## Head
         self.det_head = nn.Sequential(
             build_det_head(cfg['det_head'], self.fpn_dims, self.head_dim, self.num_levels),
             build_det_pred(self.head_dim, self.head_dim, self.strides, num_classes, 4, self.num_levels)

models/detectors/rtcdet/rtcdet_backbone.py

@@ -7,20 +7,14 @@ except:
     from rtcdet_basic import Conv, RTCBlock
 
 
-# Pretrained weights
+# MIM-Pretrained weights
 model_urls = {
-    # ImageNet-1K pretrained weight
-    "rtcnet_n": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_nano.pth",
-    "rtcnet_s": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_small.pth",
+    "rtcnet_n": None,
+    "rtcnet_t": None,
+    "rtcnet_s": None,
     "rtcnet_m": None,
     "rtcnet_l": None,
     "rtcnet_x": None,
-    # MIM-pretrained weights
-    "mae_rtcnet_n": None,
-    "mae_rtcnet_s": None,
-    "mae_rtcnet_m": None,
-    "mae_rtcnet_l": None,
-    "mae_rtcnet_x": None,
 }
 
 
@@ -109,19 +103,18 @@ def build_backbone(cfg, pretrained=False):
 
     # Model name
     width, depth, ratio = cfg['width'], cfg['depth'], cfg['ratio']
-    model_name = "{}" if not cfg['bk_pretrained_mae'] else "mae_{}"
     if  width == 0.25   and depth == 0.34 and ratio == 2.0:
-        model_name = model_name.format("rtcnet_n")
+        model_name = "rtcnet_n"
     elif width == 0.375 and depth == 0.34 and ratio == 2.0:
-        model_name = model_name.format("rtcnet_t")
+        model_name = "rtcnet_t"
     elif width == 0.50  and depth == 0.34 and ratio == 2.0:
-        model_name = model_name.format("rtcnet_s")
+        model_name = "rtcnet_s"
     elif width == 0.75  and depth == 0.67 and ratio == 1.5:
-        model_name = model_name.format("rtcnet_m")
+        model_name = "rtcnet_m"
     elif width == 1.0   and depth == 1.0  and ratio == 1.0:
-        model_name = model_name.format("rtcnet_l")
+        model_name = "rtcnet_l"
     elif width == 1.25  and depth == 1.34  and ratio == 1.0:
-        model_name = model_name.format("rtcnet_x")
+        model_name = "rtcnet_x"
     else:
         raise NotImplementedError("No such model size : width={}, depth={}, ratio={}. ".format(width, depth, ratio))