modify yolox2 basic module

models/detectors/yolox2/yolox2_backbone.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 
 try:
-    from .yolox2_basic import Conv, YoloStageBlock
+    from .yolox2_basic import Conv, Yolox2StageBlock
 except:
-    from yolox2_basic import Conv, YoloStageBlock
+    from yolox2_basic import Conv, Yolox2StageBlock
 
 
 # ---------------------------- Backbone ----------------------------
@@ -17,46 +17,46 @@ class Yolox2Backbone(nn.Module):
         # P2/4
         self.layer_2 = nn.Sequential(
             Conv(self.feat_dims[0], self.feat_dims[1], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
-            YoloStageBlock(in_dim     = self.feat_dims[1],
-                           out_dim    = self.feat_dims[1],
-                           num_blocks = round(3*depth),
-                           shortcut   = True,
-                           act_type   = act_type,
-                           norm_type  = norm_type,
-                           depthwise  = depthwise)
+            Yolox2StageBlock(in_dim     = self.feat_dims[1],
+                             out_dim    = self.feat_dims[1],
+                             num_blocks = round(3*depth),
+                             shortcut   = True,
+                             act_type   = act_type,
+                             norm_type  = norm_type,
+                             depthwise  = depthwise)
         )
         # P3/8
         self.layer_3 = nn.Sequential(
             Conv(self.feat_dims[1], self.feat_dims[2], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
-            YoloStageBlock(in_dim     = self.feat_dims[2],
-                           out_dim    = self.feat_dims[2],
-                           num_blocks = round(9*depth),
-                           shortcut   = True,
-                           act_type   = act_type,
-                           norm_type  = norm_type,
-                           depthwise  = depthwise)
+            Yolox2StageBlock(in_dim     = self.feat_dims[2],
+                             out_dim    = self.feat_dims[2],
+                             num_blocks = round(9*depth),
+                             shortcut   = True,
+                             act_type   = act_type,
+                             norm_type  = norm_type,
+                             depthwise  = depthwise)
         )
         # P4/16
         self.layer_4 = nn.Sequential(
             Conv(self.feat_dims[2], self.feat_dims[3], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
-            YoloStageBlock(in_dim     = self.feat_dims[3],
-                           out_dim    = self.feat_dims[3],
-                           num_blocks = round(9*depth),
-                           shortcut   = True,
-                           act_type   = act_type,
-                           norm_type  = norm_type,
-                           depthwise  = depthwise)
+            Yolox2StageBlock(in_dim     = self.feat_dims[3],
+                             out_dim    = self.feat_dims[3],
+                             num_blocks = round(9*depth),
+                             shortcut   = True,
+                             act_type   = act_type,
+                             norm_type  = norm_type,
+                             depthwise  = depthwise)
         )
         # P5/32
         self.layer_5 = nn.Sequential(
             Conv(self.feat_dims[3], self.feat_dims[4], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
-            YoloStageBlock(in_dim     = self.feat_dims[4],
-                           out_dim    = self.feat_dims[4],
-                           num_blocks = round(3*depth),
-                           shortcut   = True,
-                           act_type   = act_type,
-                           norm_type  = norm_type,
-                           depthwise  = depthwise)
+            Yolox2StageBlock(in_dim     = self.feat_dims[4],
+                             out_dim    = self.feat_dims[4],
+                             num_blocks = round(3*depth),
+                             shortcut   = True,
+                             act_type   = act_type,
+                             norm_type  = norm_type,
+                             depthwise  = depthwise)
         )
 
     def forward(self, x):

models/detectors/yolox2/yolox2_basic.py

@@ -105,35 +105,39 @@ class YoloBottleneck(nn.Module):
         return x + h if self.shortcut else h
 
 ## Yolo StageBlock
-class YoloStageBlock(nn.Module):
+class Yolox2StageBlock(nn.Module):
     def __init__(self,
-                 in_dim,
-                 out_dim,
-                 num_blocks = 1,
-                 shortcut   = False,
-                 act_type   = 'silu',
-                 norm_type  = 'BN',
-                 depthwise  = False,):
-        super(YoloStageBlock, self).__init__()
+                 in_dim     :int,
+                 out_dim    :int,
+                 num_blocks :int  = 1,
+                 shortcut   :bool = False,
+                 act_type   :str  = 'silu',
+                 norm_type  :str  = 'BN',
+                 depthwise  :bool = False,):
+        super(Yolox2StageBlock, self).__init__()
         self.inter_dim = out_dim // 2
         self.cv1 = Conv(in_dim, self.inter_dim, k=1, act_type=act_type, norm_type=norm_type)
         self.cv2 = Conv(in_dim, self.inter_dim, k=1, act_type=act_type, norm_type=norm_type)
-        self.m = nn.Sequential(*(
+        self.blocks = nn.Sequential(*(
             YoloBottleneck(self.inter_dim, self.inter_dim, 1.0, [1, 3], shortcut, act_type, norm_type, depthwise)
             for _ in range(num_blocks)))
-        self.output_proj = Conv((2 + num_blocks) * self.inter_dim, out_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv3 = Conv(self.inter_dim * num_blocks, self.inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.output_proj = Conv(2 * self.inter_dim, out_dim, k=1, act_type=act_type, norm_type=norm_type)
 
     def forward(self, x):
         # Input proj
         x1 = self.cv1(x)
         x2 = self.cv2(x)
-        out = list([x1, x2])
 
         # Bottleneck
-        out.extend(m(out[-1]) for m in self.m)
+        out = []
+        for m in self.blocks:
+            x2 = m(x2)
+            out.append(x2)
+        x2 = self.cv3(torch.cat(out, dim=1))
 
         # Output proj
-        out = self.output_proj(torch.cat(out, dim=1))
+        out = self.output_proj(torch.cat([x1, x2], dim=1))
 
         return out
     

models/detectors/yolox2/yolox2_pafpn.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 try:
-    from .yolox2_basic import Conv, YoloStageBlock
+    from .yolox2_basic import Conv, Yolox2StageBlock
 except:
-    from yolox2_basic import Conv, YoloStageBlock
+    from yolox2_basic import Conv, Yolox2StageBlock
 
 
 # PaFPN-ELAN
@@ -29,45 +29,45 @@ class Yolox2PaFPN(nn.Module):
         # ---------------- Top dwon ----------------
         ## P5 -> P4
         self.reduce_layer_1 = Conv(c5, round(512*width), k=1, act_type=act_type, norm_type=norm_type)
-        self.top_down_layer_1 = YoloStageBlock(in_dim       = round(512*width) + c4,
-                                               out_dim      = round(512*width),
-                                               num_blocks   = round(3*depth),
-                                               shortcut     = False,
-                                               act_type     = act_type,
-                                               norm_type    = norm_type,
-                                               depthwise    = depthwise,
-                                               )
+        self.top_down_layer_1 = Yolox2StageBlock(in_dim       = round(512*width) + c4,
+                                                 out_dim      = round(512*width),
+                                                 num_blocks   = round(3*depth),
+                                                 shortcut     = False,
+                                                 act_type     = act_type,
+                                                 norm_type    = norm_type,
+                                                 depthwise    = depthwise,
+                                                 )
         ## P4 -> P3
         self.reduce_layer_2 = Conv(round(512*width), round(256*width), k=1, act_type=act_type, norm_type=norm_type)
-        self.top_down_layer_2 = YoloStageBlock(in_dim       = round(256*width) + c3,
-                                               out_dim      = round(256*width),
-                                               num_blocks   = round(3*depth),
-                                               shortcut     = False,
-                                               act_type     = act_type,
-                                               norm_type    = norm_type,
-                                               depthwise    = depthwise,
-                                               )
+        self.top_down_layer_2 = Yolox2StageBlock(in_dim       = round(256*width) + c3,
+                                                 out_dim      = round(256*width),
+                                                 num_blocks   = round(3*depth),
+                                                 shortcut     = False,
+                                                 act_type     = act_type,
+                                                 norm_type    = norm_type,
+                                                 depthwise    = depthwise,
+                                                 )
         # ---------------- Bottom up ----------------
         ## P3 -> P4
         self.downsample_layer_1 = Conv(round(256*width), round(256*width), k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        self.bottom_up_layer_1 = YoloStageBlock(in_dim       = round(256*width) + round(256*width),
-                                                out_dim      = round(512*width),
-                                                num_blocks   = round(3*depth),
-                                                shortcut     = False,
-                                                act_type     = act_type,
-                                                norm_type    = norm_type,
-                                                depthwise    = depthwise,
-                                                )
+        self.bottom_up_layer_1 = Yolox2StageBlock(in_dim       = round(256*width) + round(256*width),
+                                                  out_dim      = round(512*width),
+                                                  num_blocks   = round(3*depth),
+                                                  shortcut     = False,
+                                                  act_type     = act_type,
+                                                  norm_type    = norm_type,
+                                                  depthwise    = depthwise,
+                                                  )
         ## P4 -> P5
         self.downsample_layer_2 = Conv(round(512*width), round(512*width), k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        self.bottom_up_layer_2 = YoloStageBlock(in_dim       = round(512 * width) + round(512 * width),
-                                                out_dim      = round(1024 * width),
-                                                num_blocks   = round(3*depth),
-                                                shortcut     = False,
-                                                act_type     = act_type,
-                                                norm_type    = norm_type,
-                                                depthwise    = depthwise,
-                                                )
+        self.bottom_up_layer_2 = Yolox2StageBlock(in_dim       = round(512 * width) + round(512 * width),
+                                                  out_dim      = round(1024 * width),
+                                                  num_blocks   = round(3*depth),
+                                                  shortcut     = False,
+                                                  act_type     = act_type,
+                                                  norm_type    = norm_type,
+                                                  depthwise    = depthwise,
+                                                  )
         ## output proj layers
         if out_dim is not None:
             self.out_layers = nn.ModuleList([