decoupled_head_ECA

models/yolo.py

@@ -126,6 +126,23 @@ class Detect(nn.Module):
         anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
         return grid, anchor_grid
     
+
+# 新增
+class ECABlock(nn.Module):
+    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):
+        y = self.avg_pool(x)
+        y = y.squeeze(-1).transpose(-1, -2)
+        y = self.conv(y)
+        y = self.sigmoid(y)
+        y = y.transpose(-1, -2).unsqueeze(-1)
+        return x * y.expand_as(x)
+
 # 新增
 class DetectDecoupledLite(nn.Module):
     stride = None
@@ -199,6 +216,76 @@ class DetectDecoupledLite(nn.Module):
         grid = torch.stack((xv, yv), 2).expand(shape) - 0.5
         anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
         return grid, anchor_grid
+    
+# 新增
+class DetectDecoupledLiteECA(nn.Module):
+    stride = None
+    dynamic = False
+    export = False
+
+    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
+        super().__init__()
+        self.nc = nc
+        self.no = nc + 5
+        self.nl = len(anchors)
+        self.na = len(anchors[0]) // 2
+        self.grid = [torch.empty(0) for _ in range(self.nl)]
+        self.anchor_grid = [torch.empty(0) for _ in range(self.nl)]
+        self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2))
+        self.inplace = inplace
+
+        head_c = 128
+
+        self.stems = nn.ModuleList(Conv(c, head_c, 1, 1) for c in ch)
+        self.regobj_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)
+        self.cls_attn = nn.ModuleList(ECABlock(head_c, k_size=3) for _ in ch)
+
+        self.box_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * 4, 1) for _ in ch)
+        self.obj_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * 1, 1) for _ in ch)
+        self.cls_preds = nn.ModuleList(nn.Conv2d(head_c, self.na * self.nc, 1) for _ in ch)
+
+    def forward(self, x):
+        z = []
+        for i in range(self.nl):
+            xi = self.stems[i](x[i])
+
+            regobj_feat = self.regobj_convs[i](xi)
+
+            cls_feat = self.cls_convs[i](xi)
+            cls_feat = self.cls_attn[i](cls_feat)
+
+            box = self.box_preds[i](regobj_feat)
+            obj = self.obj_preds[i](regobj_feat)
+            cls = self.cls_preds[i](cls_feat)
+
+            out = torch.cat((box, obj, cls), 1)
+
+            bs, _, ny, nx = out.shape
+            out = out.view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
+            x[i] = out
+
+            if not self.training:
+                if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
+                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
+
+                xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
+                xy = (xy * 2 + self.grid[i]) * self.stride[i]
+                wh = (wh * 2) ** 2 * self.anchor_grid[i]
+                y = torch.cat((xy, wh, conf), 4)
+                z.append(y.view(bs, self.na * nx * ny, self.no))
+
+        return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
+
+    def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")):
+        d = self.anchors[i].device
+        t = self.anchors[i].dtype
+        shape = 1, self.na, ny, nx, 2
+        y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
+        yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x)
+        grid = torch.stack((xv, yv), 2).expand(shape) - 0.5
+        anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
+        return grid, anchor_grid
 
 
 class Segment(Detect):
@@ -282,10 +369,11 @@ class BaseModel(nn.Module):
         self = super()._apply(fn)
         m = self.model[-1]  # Detect()
         # if isinstance(m, (Detect, Segment)):
-        if isinstance(m, (Detect, DetectDecoupledLite, Segment)):
+        # if isinstance(m, (Detect, DetectDecoupledLite, Segment)):
+        if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteECA, Segment)):
             m.stride = fn(m.stride)
             m.grid = list(map(fn, m.grid))
-            if isinstance(m.anchor_grid, list):
+            if isinstance(m.anchor_grid, list): 
                 m.anchor_grid = list(map(fn, m.anchor_grid))
         return self
 
@@ -320,7 +408,8 @@ class DetectionModel(BaseModel):
         # Build strides, anchors
         m = self.model[-1]  # Detect()
         # if isinstance(m, (Detect, Segment)):
-        if isinstance(m, (Detect, DetectDecoupledLite, Segment)):
+        # if isinstance(m, (Detect, DetectDecoupledLite, Segment)):
+        if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteECA, Segment)):
 
             def _forward(x):
                 """Passes the input 'x' through the model and returns the processed output."""
@@ -337,6 +426,8 @@ class DetectionModel(BaseModel):
                 self._initialize_biases()
             elif isinstance(m, DetectDecoupledLite):
                 self._initialize_decoupled_biases()
+            elif isinstance(m, DetectDecoupledLiteECA):
+                self._initialize_decoupled_eca_biases()
 
         # Init weights, biases
         initialize_weights(self)
@@ -363,6 +454,24 @@ class DetectionModel(BaseModel):
                 b.data += torch.log(cf / cf.sum())
             mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
 
+    # 新增
+    def _initialize_decoupled_eca_biases(self, cf=None):
+        """Initialize biases for DetectDecoupledLiteECA."""
+        m = self.model[-1]
+
+        for mi, s in zip(m.obj_preds, m.stride):
+            b = mi.bias.view(m.na, -1)
+            b.data += math.log(8 / (640 / s) ** 2)
+            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+
+        for mi in m.cls_preds:
+            b = mi.bias.view(m.na, -1)
+            if cf is None:
+                b.data += math.log(0.6 / (m.nc - 0.99999))
+            else:
+                b.data += torch.log(cf / cf.sum())
+            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+
     def forward(self, x, augment=False, profile=False, visualize=False):
         """Performs single-scale or augmented inference and may include profiling or visualization."""
         if augment:
@@ -535,7 +644,8 @@ def parse_model(d, ch):
             c2 = sum(ch[x] for x in f)
         # TODO: channel, gw, gd
         # elif m in {Detect, Segment}:
-        elif m in {Detect, DetectDecoupledLite, Segment}:
+        # elif m in {Detect, DetectDecoupledLite, Segment}:
+        elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteECA, Segment}:
             args.append([ch[x] for x in f])
             if isinstance(args[1], int):  # number of anchors
                 args[1] = [list(range(args[1] * 2))] * len(f)