cls_head_enhance_Lightweight

models/yolo.py

@@ -125,7 +125,19 @@ class Detect(nn.Module):
         grid = torch.stack((xv, yv), 2).expand(shape) - 0.5  # add grid offset, i.e. y = 2.0 * x - 0.5
         anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape)
         return grid, anchor_grid
-    
+
+# 新增
+class ClsEnhanceBlock(nn.Module):
+    """Lightweight classification enhancement block for G11."""
+    def __init__(self, c1, expand=0.5):
+        super().__init__()
+        c_ = max(16, int(c1 * expand))
+        self.cv1 = Conv(c1, c_, 3, 1)
+        self.cv2 = Conv(c_, c1, 3, 1)
+
+    def forward(self, x):
+        return x + self.cv2(self.cv1(x))
+
 # 新增
 class DetectDecoupledLite(nn.Module):
     stride = None
@@ -199,6 +211,84 @@ 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 DetectDecoupledLiteClsEnhance(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
+
+        # stem
+        self.stems = nn.ModuleList(Conv(c, head_c, 1, 1) for c in ch)
+
+        # reg + obj shared trunk
+        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)
+
+        # G11: cls enhance block
+        self.cls_enhance = nn.ModuleList(ClsEnhanceBlock(head_c, expand=0.5) for _ in ch)
+
+        # predictors
+        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_enhance[i](cls_feat)   # G11 核心增强
+
+            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,7 +372,8 @@ 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, DetectDecoupledLiteClsEnhance, Segment)):
             m.stride = fn(m.stride)
             m.grid = list(map(fn, m.grid))
             if isinstance(m.anchor_grid, list):
@@ -320,7 +411,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, DetectDecoupledLiteClsEnhance, Segment)):
 
             def _forward(x):
                 """Passes the input 'x' through the model and returns the processed output."""
@@ -337,6 +429,8 @@ class DetectionModel(BaseModel):
                 self._initialize_biases()
             elif isinstance(m, DetectDecoupledLite):
                 self._initialize_decoupled_biases()
+            elif isinstance(m, DetectDecoupledLiteClsEnhance):
+                self._initialize_decoupled_cls_enhance_biases()
 
         # Init weights, biases
         initialize_weights(self)
@@ -363,6 +457,26 @@ class DetectionModel(BaseModel):
                 b.data += torch.log(cf / cf.sum())
             mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
 
+    # 新增
+    def _initialize_decoupled_cls_enhance_biases(self, cf=None):
+        """Initialize biases for DetectDecoupledLiteClsEnhance."""
+        m = self.model[-1]
+
+        # objectness bias
+        for mi, s in zip(m.obj_preds, m.stride):
+            b = mi.bias.view(m.na, -1)  # (na, 1)
+            b.data += math.log(8 / (640 / s) ** 2)
+            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+
+        # classification bias
+        for mi in m.cls_preds:
+            b = mi.bias.view(m.na, -1)  # (na, nc)
+            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 +649,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, DetectDecoupledLiteClsEnhance, 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)