[G1][Backbone] add CBAM and export scripts

models/common.py

@@ -1121,3 +1121,58 @@ 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 ChannelAttention(nn.Module):
+    def __init__(self, channels, reduction=16):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc = nn.Sequential(
+            nn.Conv2d(channels, channels // reduction, 1, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channels // reduction, channels, 1, bias=False)
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(self.max_pool(x))
+        return self.sigmoid(avg_out + max_out)
+
+class SpatialAttention(nn.Module):
+    def __init__(self, kernel_size=7):
+        super().__init__()
+        assert kernel_size in (3, 7)
+        padding = 3 if kernel_size == 7 else 1
+
+        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        return self.sigmoid(self.conv(x))
+
+class CBAM(nn.Module):
+    def __init__(self, channels, reduction=16, kernel_size=7):
+        super().__init__()
+        self.ca = ChannelAttention(channels, reduction)
+        self.sa = SpatialAttention(kernel_size)
+
+    def forward(self, x):
+        x = x * self.ca(x)
+        x = x * self.sa(x)
+        return x
+    
+class C3CBAM(C3):  
+    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
+        super().__init__(c1, c2, n, shortcut, g, e)
+        self.cbam = CBAM(c2)
+
+    def forward(self, x):
+        x = super().forward(x)
+        return self.cbam(x)

models/yolo.py

@@ -27,6 +27,8 @@ if platform.system() != "Windows":
 
 from models.common import (
     C3,
+    C3CBAM,     # ★ 新增
+    CBAM,       # ★ 可选（严格来说不是必须，但建议）
     C3SPP,
     C3TR,
     SPP,
@@ -414,6 +416,7 @@ def parse_model(d, ch):
             CrossConv,
             BottleneckCSP,
             C3,
+            C3CBAM,
             C3TR,
             C3SPP,
             C3Ghost,
@@ -426,7 +429,7 @@ def parse_model(d, ch):
                 c2 = make_divisible(c2 * gw, ch_mul)
 
             args = [c1, c2, *args[1:]]
-            if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
+            if m in {BottleneckCSP, C3, C3CBAM, C3TR, C3Ghost, C3x}:
                 args.insert(2, n)  # number of repeats
                 n = 1
         elif m is nn.BatchNorm2d:

models/yolov5s.yaml

@@ -10,19 +10,32 @@ anchors:
   - [116, 90, 156, 198, 373, 326] # P5/32
 
 # YOLOv5 v6.0 backbone
+# backbone:
+#   # [from, number, module, args]
+#   [
+#     [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
+#     [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
+#     [-1, 3, C3, [128]],
+#     [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
+#     [-1, 6, C3, [256]],
+#     [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
+#     [-1, 9, C3, [512]],
+#     [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
+#     [-1, 3, C3, [1024]],
+#     [-1, 1, SPPF, [1024, 5]], # 9
+#   ]
 backbone:
-  # [from, number, module, args]
   [
-    [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
-    [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
-    [-1, 3, C3, [128]],
-    [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
-    [-1, 6, C3, [256]],
-    [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
-    [-1, 9, C3, [512]],
-    [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
-    [-1, 3, C3, [1024]],
-    [-1, 1, SPPF, [1024, 5]], # 9
+    [-1, 1, Conv, [64, 6, 2, 2]],          # 0-P1/2
+    [-1, 1, Conv, [128, 3, 2]],            # 1-P2/4
+    [-1, 3, C3, [128]],                    # 2 (保持不变)
+    [-1, 1, Conv, [256, 3, 2]],            # 3-P3/8
+    [-1, 6, C3, [256]],                    # 4 (保持不变)
+    [-1, 1, Conv, [512, 3, 2]],            # 5-P4/16
+    [-1, 9, C3CBAM, [512]],                # 6 ★ 加 CBAM
+    [-1, 1, Conv, [1024, 3, 2]],           # 7-P5/32
+    [-1, 3, C3CBAM, [1024]],               # 8 ★ 加 CBAM
+    [-1, 1, SPPF, [1024, 5]],               # 9
   ]
 
 # YOLOv5 v6.0 head