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