yjh0410 преди 2 години
родител
ревизия
63a7b3f516
променени са 2 файла, в които са добавени 64 реда и са изтрити 2 реда
  1. 59 0
      config/model_config/rtcdet_v2_config.py
  2. 5 2
      models/detectors/rtcdet_v2/rtcdet_v2_pafpn.py

+ 59 - 0
config/model_config/rtcdet_v2_config.py

@@ -2,6 +2,65 @@
 
 
 rtcdet_v2_cfg = {
+    'rtcdet_v2_s':{
+        # ---------------- Model config ----------------
+        ## Backbone
+        'backbone': 'elannet_v2',
+        'pretrained': True,
+        'bk_act': 'silu',
+        'bk_norm': 'BN',
+        'bk_depthwise': False,
+        'width': 0.50,
+        'depth': 0.34,
+        '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_reduce_layer': 'conv',
+        'fpn_downsample_layer': 'conv',
+        'fpn_core_block': 'elan_block',
+        'fpn_expand_ratio': 0.25,
+        '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,
+        'reg_max': 16,
+        # ---------------- Train config ----------------
+        ## Input
+        'multi_scale': [0.5, 1.25],   # 320 -> 800
+        'trans_type': 'yolox_small',
+        # ---------------- Assignment config ----------------
+        ## Matcher
+        'matcher': {'tal': {'topk': 10,
+                            'alpha': 0.5,
+                            'beta': 6.0},
+                    'ota': {'center_sampling_radius': 2.5,
+                             'topk_candidate': 10},
+                    },
+        # ---------------- Loss config ----------------
+        ## Loss weight
+        'ema_update': False,
+        'loss_box_aux': True,
+        'loss_cls_weight': {'tal': 0.5, 'ota': 1.0},
+        'loss_box_weight': {'tal': 7.0, 'ota': 5.0},
+        'loss_dfl_weight': {'tal': 1.5, 'ota': 1.0},
+        # ---------------- Train config ----------------
+        'trainer_type': 'yolox',
+    },
+
     'rtcdet_v2_l':{
         # ---------------- Model config ----------------
         ## Backbone

+ 5 - 2
models/detectors/rtcdet_v2/rtcdet_v2_pafpn.py

@@ -17,8 +17,11 @@ class RTCDetPaFPN(nn.Module):
         self.fpn_dims = [round(256*cfg['width']), round(512*cfg['width']), round(1024*cfg['width'])]
 
         # --------------------------- Input proj ---------------------------
-        self.input_projs = nn.ModuleList([nn.Conv2d(in_dim, fpn_dim, kernel_size=1)
-                                          for in_dim, fpn_dim in zip(in_dims, self.fpn_dims)])
+        if in_dims == self.fpn_dims:
+            self.input_projs = nn.ModuleList([nn.Identity() for _ in range(len(in_dims))])
+        else:
+            self.input_projs = nn.ModuleList([nn.Conv2d(in_dim, fpn_dim, kernel_size=1)
+                                            for in_dim, fpn_dim in zip(in_dims, self.fpn_dims)])
                 
         # --------------------------- Top-down FPN ---------------------------
         ## P5 -> P4