add decoupled head

export_my_gra.sh

@@ -1,6 +1,6 @@
 python export.py \
-  --weights /data6/zhangw/ZW_gra_project/yolov5_result_pretrained/G0_cub100_yolov5s6_raw_e300/weights/best.pt \
-  --img 640 640 \
+  --weights /data6/zhangw/yolov5/yolov5s6.pt \
+  --img 640 \
   --batch 1 \
   --device cpu \
   --include onnx \

models/yolo.py

@@ -126,6 +126,90 @@ 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 DecoupledDetect(nn.Module):
+    """YOLOv5 decoupled detect head: anchor-based, training/inference interface kept the same as Detect."""
+    stride = None
+    dynamic = False
+    export = False
+
+    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):
+        super().__init__()
+        self.nc = nc                    # number of classes
+        self.no = nc + 5               # outputs per anchor = [x, y, w, h, obj] + cls
+        self.nl = len(anchors)         # number of detection layers
+        self.na = len(anchors[0]) // 2 # number of anchors
+        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
+
+        # hidden channels for each scale
+        c_ = [max(128, x // 2) for x in ch]
+
+        # stem
+        self.stems = nn.ModuleList(Conv(x, c, 1, 1) for x, c in zip(ch, c_))
+
+        # classification branch
+        self.cls_convs = nn.ModuleList(
+            nn.Sequential(
+                Conv(c, c, 3, 1),
+                Conv(c, c, 3, 1),
+            ) for c in c_
+        )
+
+        # regression/objectness branch
+        self.reg_convs = nn.ModuleList(
+            nn.Sequential(
+                Conv(c, c, 3, 1),
+                Conv(c, c, 3, 1),
+            ) for c in c_
+        )
+
+        # prediction layers
+        self.cls_preds = nn.ModuleList(nn.Conv2d(c, self.na * self.nc, 1) for c in c_)
+        self.reg_preds = nn.ModuleList(nn.Conv2d(c, self.na * 4, 1) for c in c_)
+        self.obj_preds = nn.ModuleList(nn.Conv2d(c, self.na * 1, 1) for c in c_)
+
+    def forward(self, x):
+        z = []  # inference output
+        for i in range(self.nl):
+            xi = self.stems[i](x[i])
+
+            cls_feat = self.cls_convs[i](xi)
+            reg_feat = self.reg_convs[i](xi)
+
+            cls_out = self.cls_preds[i](cls_feat)   # [bs, na*nc, ny, nx]
+            reg_out = self.reg_preds[i](reg_feat)   # [bs, na*4,  ny, nx]
+            obj_out = self.obj_preds[i](reg_feat)   # [bs, na*1,  ny, nx]
+
+            # keep original YOLOv5 output order: [xywh, obj, cls]
+            x[i] = torch.cat((reg_out, obj_out, cls_out), 1)
+
+            bs, _, ny, nx = x[i].shape
+            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
+
+            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):
     """YOLOv5 Segment head for segmentation models, extending Detect with mask and prototype layers."""
@@ -207,7 +291,8 @@ class BaseModel(nn.Module):
         """
         self = super()._apply(fn)
         m = self.model[-1]  # Detect()
-        if isinstance(m, (Detect, Segment)):
+        # if isinstance(m, (Detect, Segment)):
+        if isinstance(m, (Detect, DecoupledDetect, Segment)):
             m.stride = fn(m.stride)
             m.grid = list(map(fn, m.grid))
             if isinstance(m.anchor_grid, list):
@@ -244,7 +329,8 @@ class DetectionModel(BaseModel):
 
         # Build strides, anchors
         m = self.model[-1]  # Detect()
-        if isinstance(m, (Detect, Segment)):
+        # if isinstance(m, (Detect, Segment)):
+        if isinstance(m, (Detect, DecoupledDetect, Segment)):
 
             def _forward(x):
                 """Passes the input 'x' through the model and returns the processed output."""
@@ -314,26 +400,51 @@ class DetectionModel(BaseModel):
         y[-1] = y[-1][:, i:]  # small
         return y
 
-    def _initialize_biases(self, cf=None):
-        """
-        Initializes biases for YOLOv5's Detect() module, optionally using class frequencies (cf).
+    # def _initialize_biases(self, cf=None):
+    #     """
+    #     Initializes biases for YOLOv5's Detect() module, optionally using class frequencies (cf).
 
-        For details see https://arxiv.org/abs/1708.02002 section 3.3.
-        """
+    #     For details see https://arxiv.org/abs/1708.02002 section 3.3.
+    #     """
+    #     # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
+    #     m = self.model[-1]  # Detect() module
+    #     for mi, s in zip(m.m, m.stride):  # from
+    #         b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
+    #         b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
+    #         b.data[:, 5 : 5 + m.nc] += (
+    #             math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
+    #         )  # cls
+    #         mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+    def _initialize_biases(self, cf=None):
+        """Initialize biases for Detect() or DecoupledDetect()."""
         # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
         m = self.model[-1]  # Detect() module
-        for mi, s in zip(m.m, m.stride):  # from
-            b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
-            b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
-            b.data[:, 5 : 5 + m.nc] += (
-                math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
-            )  # cls
-            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
 
+        if isinstance(m, Detect):
+            for mi, s in zip(m.m, m.stride):
+                b = mi.bias.view(m.na, -1)  # conv.bias(255) to (3,85)
+                b.data[:, 4] += math.log(8 / (640 / s) ** 2)  # obj
+                b.data[:, 5:5 + m.nc] += (
+                    math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
+                )  # cls
+                mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
+
+        elif isinstance(m, DecoupledDetect):
+            for cls_pred, obj_pred, s in zip(m.cls_preds, m.obj_preds, m.stride):
+                # objectness bias
+                b_obj = obj_pred.bias.view(m.na, -1)  # [na, 1]
+                b_obj.data += math.log(8 / (640 / s) ** 2)
+                obj_pred.bias = torch.nn.Parameter(b_obj.view(-1), requires_grad=True)
+
+                # classification bias
+                b_cls = cls_pred.bias.view(m.na, -1)  # [na, nc]
+                b_cls.data += (
+                    math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())
+                )
+                cls_pred.bias = torch.nn.Parameter(b_cls.view(-1), requires_grad=True)
 
 Model = DetectionModel  # retain YOLOv5 'Model' class for backwards compatibility
 
-
 class SegmentationModel(DetectionModel):
     """YOLOv5 segmentation model for object detection and segmentation tasks with configurable parameters."""
 
@@ -434,7 +545,8 @@ def parse_model(d, ch):
         elif m is Concat:
             c2 = sum(ch[x] for x in f)
         # TODO: channel, gw, gd
-        elif m in {Detect, Segment}:
+        # elif m in {Detect, Segment}:
+        elif m in {Detect, DecoupledDetect, 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)

train_my_gra.sh

@@ -1,9 +1,10 @@
-export CUDA_VISIBLE_DEVICES=4,5,6,7
+export CUDA_VISIBLE_DEVICES=0,1,2,3
 python /data6/zhangw/yolov5/train.py \
     --img 960 \
     --batch 16 \
     --epochs 300 \
     --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \
     --weights /data6/zhangw/yolov5/yolov5s6.pt \
-    --project /data6/zhangw/ZW_gra_project/yolov5_result \
-    --name G0_cub100_yolov5s6_raw_e300
+    --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \
+    --name G6_cub100_yolov5s6_decoupled_head_e300 \
+    --cfg /data6/zhangw/yolov5/models/hub/yolov5s6_decoupled.yaml