decoupled_light_FGEM

models/common.py

@@ -1121,3 +1121,57 @@ 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, c1, 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.expand_as(x)
+
+
+class FGEMLite(nn.Module):
+    """Lightweight Fine-Grained Enhancement Module for classification branch only."""
+    def __init__(self, c1, c_mid=None):
+        super().__init__()
+        c_mid = c_mid or max(c1 // 2, 32)
+
+        # 1x1 reduce
+        self.reduce = Conv(c1, c_mid, 1, 1)
+
+        # depthwise local modeling
+        self.dw = nn.Sequential(
+            nn.Conv2d(c_mid, c_mid, kernel_size=3, stride=1, padding=1, groups=c_mid, bias=False),
+            nn.BatchNorm2d(c_mid),
+            nn.SiLU()
+        )
+
+        # pointwise fusion
+        self.pw = Conv(c_mid, c_mid, 1, 1)
+
+        # channel attention
+        self.eca = ECA(c_mid, k_size=3)
+
+        # 1x1 expand
+        self.expand = Conv(c_mid, c1, 1, 1)
+
+    def forward(self, x):
+        identity = x
+        y = self.reduce(x)
+        y = self.dw(y)
+        y = self.pw(y)
+        y = self.eca(y)
+        y = self.expand(y)
+        return identity + y

models/yolo.py

@@ -48,6 +48,7 @@ from models.common import (
     GhostBottleneck,
     GhostConv,
     Proto,
+    FGEMLite,   # 新增
 )
 from models.experimental import MixConv2d
 from utils.autoanchor import check_anchor_order
@@ -152,7 +153,14 @@ class DetectDecoupledLite(nn.Module):
         self.regobj_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch)
 
         # cls trunk
-        self.cls_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch)
+        # self.cls_convs = nn.ModuleList(Conv(head_c, head_c, 3, 1) for _ in ch)
+        # cls trunk + FGEMLite
+        self.cls_convs = nn.ModuleList(
+            nn.Sequential(
+                Conv(head_c, head_c, 3, 1),
+                FGEMLite(head_c, head_c // 2)
+            ) for _ in ch
+        )
 
         # predictors
         self.box_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * 4, 1) for _ in ch)

models/yolov5s6_decoupled_light_FGEM.yaml

@@ -0,0 +1,62 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+# Parameters
+nc: 100 # number of classes
+depth_multiple: 0.33 # model depth multiple
+width_multiple: 0.50 # layer channel multiple
+anchors:
+  - [19, 27, 44, 40, 38, 94] # P3/8
+  - [96, 68, 86, 152, 180, 137] # P4/16
+  - [140, 301, 303, 264, 238, 542] # P5/32
+  - [436, 615, 739, 380, 925, 792] # P6/64
+
+# 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, [768, 3, 2]], # 7-P5/32
+    [-1, 3, C3, [768]],
+    [-1, 1, Conv, [1024, 3, 2]], # 9-P6/64
+    [-1, 3, C3, [1024]],
+    [-1, 1, SPPF, [1024, 5]], # 11
+  ]
+
+# YOLOv5 v6.0 head
+head: [
+    [-1, 1, Conv, [768, 1, 1]],
+    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
+    [[-1, 8], 1, Concat, [1]], # cat backbone P5
+    [-1, 3, C3, [768, False]], # 15
+
+    [-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]], # 19
+
+    [-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]], # 23 (P3/8-small)
+
+    [-1, 1, Conv, [256, 3, 2]],
+    [[-1, 20], 1, Concat, [1]], # cat head P4
+    [-1, 3, C3, [512, False]], # 26 (P4/16-medium)
+
+    [-1, 1, Conv, [512, 3, 2]],
+    [[-1, 16], 1, Concat, [1]], # cat head P5
+    [-1, 3, C3, [768, False]], # 29 (P5/32-large)
+
+    [-1, 1, Conv, [768, 3, 2]],
+    [[-1, 12], 1, Concat, [1]], # cat head P6
+    [-1, 3, C3, [1024, False]], # 32 (P6/64-xlarge)
+
+    # [[23, 26, 29, 32], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5, P6)
+    [[23, 26, 29, 32], 1, DetectDecoupledLite, [nc, anchors]], # Detect(P3, P4, P5, P6)
+  ]

train_my_gra.sh

@@ -6,5 +6,5 @@ python /data6/zhangw/yolov5/train.py \
     --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \
     --weights /data6/zhangw/yolov5/yolov5s6.pt \
     --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \
-    --name G8_decoupled_head_light_e300 \
+    --name G9_decoupled_head_light_FGEM_e300 \
-    --cfg /data6/zhangw/yolov5/models/yolov5s6_decoupled_light.yaml
+    --cfg /data6/zhangw/yolov5/models/yolov5s6_decoupled_light_FGEM.yaml