<Neck> add ECA after P3P4P5

models/common.py

@@ -1121,3 +1121,25 @@ class Classify(nn.Module):
         if isinstance(x, list):
             x = torch.cat(x, 1)
         return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
+
+# 以下是我自己添加的
+class ECA(nn.Module):
+    """Efficient Channel Attention"""
+    def __init__(self, channels, k_size=3):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv1d(
+            1, 1, kernel_size=k_size,
+            padding=(k_size - 1) // 2,
+            bias=False
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        # x: [B, C, H, W]
+        y = self.avg_pool(x)                  # [B, C, 1, 1]
+        y = y.squeeze(-1).transpose(-1, -2)   # [B, 1, C]
+        y = self.conv(y)                      # [B, 1, C]
+        y = self.sigmoid(y)
+        y = y.transpose(-1, -2).unsqueeze(-1) # [B, C, 1, 1]
+        return x * y

models/yolov5s.yaml

@@ -25,25 +25,54 @@ backbone:
     [-1, 1, SPPF, [1024, 5]], # 9
   ]
 
-# YOLOv5 v6.0 head
-head: [
-    [-1, 1, Conv, [512, 1, 1]],
-    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
-    [[-1, 6], 1, Concat, [1]], # cat backbone P4
-    [-1, 3, C3, [512, False]], # 13
+# # YOLOv5 v6.0 head
+# head: [
+#     [-1, 1, Conv, [512, 1, 1]],
+#     [-1, 1, nn.Upsample, [None, 2, "nearest"]],
+#     [[-1, 6], 1, Concat, [1]], # cat backbone P4
+#     [-1, 3, C3, [512, False]], # 13
 
-    [-1, 1, Conv, [256, 1, 1]],
-    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
-    [[-1, 4], 1, Concat, [1]], # cat backbone P3
-    [-1, 3, C3, [256, False]], # 17 (P3/8-small)
+#     [-1, 1, Conv, [256, 1, 1]],
+#     [-1, 1, nn.Upsample, [None, 2, "nearest"]],
+#     [[-1, 4], 1, Concat, [1]], # cat backbone P3
+#     [-1, 3, C3, [256, False]], # 17 (P3/8-small)
 
-    [-1, 1, Conv, [256, 3, 2]],
-    [[-1, 14], 1, Concat, [1]], # cat head P4
-    [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
+#     [-1, 1, Conv, [256, 3, 2]],
+#     [[-1, 14], 1, Concat, [1]], # cat head P4
+#     [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
 
-    [-1, 1, Conv, [512, 3, 2]],
-    [[-1, 10], 1, Concat, [1]], # cat head P5
-    [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
+#     [-1, 1, Conv, [512, 3, 2]],
+#     [[-1, 10], 1, Concat, [1]], # cat head P5
+#     [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
 
-    [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
+#     [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
+#   ]
+
+head:
+  [
+    # -------- Top-down pathway (FPN) --------
+    [-1, 1, Conv, [512, 1, 1]],                    # 10 reduce channels (P5)
+    [-1, 1, nn.Upsample, [None, 2, "nearest"]],    # 11 upsample
+    [[-1, 6], 1, Concat, [1]],                     # 12 concat backbone P4
+    [-1, 3, C3, [512, False]],                     # 13 P4 fused feature
+
+    [-1, 1, Conv, [256, 1, 1]],                    # 14 reduce channels
+    [-1, 1, nn.Upsample, [None, 2, "nearest"]],    # 15 upsample
+    [[-1, 4], 1, Concat, [1]],                     # 16 concat backbone P3
+    [-1, 3, C3, [256, False]],                     # 17 P3 fused feature
+    [-1, 1, ECA, [256]],                           # 18 ECA on P3 (semantic calibration)
+
+    # -------- Bottom-up pathway (PAN) --------
+    [-1, 1, Conv, [256, 3, 2]],                    # 19 downsample
+    [[-1, 14], 1, Concat, [1]],                    # 20 concat head P4
+    [-1, 3, C3, [512, False]],                     # 21 P4 fused feature
+    [-1, 1, ECA, [512]],                           # 22 ECA on P4
+
+    [-1, 1, Conv, [512, 3, 2]],                    # 23 downsample
+    [[-1, 10], 1, Concat, [1]],                    # 24 concat head P5
+    [-1, 3, C3, [1024, False]],                    # 25 P5 fused feature
+    [-1, 1, ECA, [1024]],                          # 26 ECA on P5
+
+    # -------- Detection heads --------
+    [[18, 22, 26], 1, Detect, [nc, anchors]],      # 27 Detect(P3, P4, P5)
   ]

train_my_gra.sh

@@ -0,0 +1,9 @@
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+python /data6/zhangw/yolov5/train.py \
+    --img 960 \
+    --batch 16 \
+    --epochs 300 \
+    --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \
+    --weights /data6/zhangw/yolov5/yolov5s6.pt \
+    --project /data6/zhangw/ZW_gra_project/yolov5_result \
+    --name G2_cub100_yolov5s6_changeNeck_e300