Browse Source

fix bug in spp_csp

yjh0410 9 months ago
parent
commit
280150e962
2 changed files with 9 additions and 10 deletions
  1. 1 1
      yolo/models/yolov7/yolov7.py
  2. 8 9
      yolo/models/yolov7/yolov7_neck.py

+ 1 - 1
yolo/models/yolov7/yolov7.py

@@ -30,7 +30,7 @@ class Yolov7(nn.Module):
         
         # ------------------- Network Structure -------------------
         self.backbone = Yolov7Backbone(use_pretrained=cfg.use_pretrained)
-        self.neck     = SPPFBlockCSP(self.backbone.feat_dims[-1], self.backbone.feat_dims[-1] // 2, expand_ratio=0.5)
+        self.neck     = SPPFBlockCSP(self.backbone.feat_dims[-1], self.backbone.feat_dims[-1] // 2, expansion=0.5)
         self.backbone.feat_dims[-1] = self.backbone.feat_dims[-1] // 2
         self.fpn      = Yolov7PaFPN(self.backbone.feat_dims[-3:], head_dim=cfg.head_dim)
         self.non_shared_heads = nn.ModuleList([DecoupledHead(cfg, in_dim)

+ 8 - 9
yolo/models/yolov7/yolov7_neck.py

@@ -11,9 +11,9 @@ class SPPF(nn.Module):
     """
         This code referenced to https://github.com/ultralytics/yolov5
     """
-    def __init__(self, in_dim, out_dim, expand_ratio=0.5):
+    def __init__(self, in_dim, out_dim, expansion=0.5):
         super().__init__()
-        inter_dim = int(in_dim * expand_ratio)
+        inter_dim = int(in_dim * expansion)
         self.out_dim = out_dim
         self.cv1 = ConvModule(in_dim, inter_dim, kernel_size=1)
         self.cv2 = ConvModule(inter_dim * 4, out_dim, kernel_size=1)
@@ -30,25 +30,24 @@ class SPPFBlockCSP(nn.Module):
     def __init__(self,
                  in_dim: int,
                  out_dim: int,
-                 expand_ratio: float = 0.5,
+                 expansion: float = 0.5,
                  ):
         super(SPPFBlockCSP, self).__init__()
-        inter_dim = int(in_dim * expand_ratio)
+        inter_dim = int(in_dim * expansion)
         self.out_dim = out_dim
         self.cv1 = ConvModule(in_dim, inter_dim, kernel_size=1)
         self.cv2 = ConvModule(in_dim, inter_dim, kernel_size=1)
         self.m = nn.Sequential(
             ConvModule(inter_dim, inter_dim, kernel_size=3),
-            SPPF(inter_dim, inter_dim, expand_ratio=1.0),
+            SPPF(inter_dim, inter_dim, expansion=1.0),
             ConvModule(inter_dim, inter_dim, kernel_size=3)
         )
         self.cv3 = ConvModule(inter_dim * 2, self.out_dim, kernel_size=1)
 
     def forward(self, x):
         x1 = self.cv1(x)
-        x2 = self.cv2(x)
-        x3 = self.m(x2)
-        y = self.cv3(torch.cat([x1, x3], dim=1))
+        x2 = self.m(self.cv2(x))
+        y = self.cv3(torch.cat([x1, x2], dim=1))
 
         return y
 
@@ -59,7 +58,7 @@ if __name__=='__main__':
     # Build a neck
     in_dim  = 512
     out_dim = 512
-    model = SPPFBlockCSP(512, 512, expand_ratio=0.5)
+    model = SPPFBlockCSP(512, 512, expansion=0.5)
 
     # Randomly generate a input data
     x = torch.randn(2, in_dim, 20, 20)