change to light decoupled head

models/yolo.py

@@ -125,6 +125,80 @@ 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 DetectDecoupledLite(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)
+
+        # 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)
+
+            box = self.box_preds[i](regobj_feat)
+            obj = self.obj_preds[i](regobj_feat)
+            cls = self.cls_preds[i](cls_feat)
+
+            # [bs, na*4 + na*1 + na*nc, ny, nx]
+            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):
@@ -207,7 +281,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, DetectDecoupledLite, Segment)):
             m.stride = fn(m.stride)
             m.grid = list(map(fn, m.grid))
             if isinstance(m.anchor_grid, list):
@@ -244,7 +319,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, DetectDecoupledLite, Segment)):
 
             def _forward(x):
                 """Passes the input 'x' through the model and returns the processed output."""
@@ -256,12 +332,36 @@ class DetectionModel(BaseModel):
             check_anchor_order(m)
             m.anchors /= m.stride.view(-1, 1, 1)
             self.stride = m.stride
-            self._initialize_biases()  # only run once
+            # self._initialize_biases()  # only run once
+            if isinstance(m, Detect):
+                self._initialize_biases()
+            elif isinstance(m, DetectDecoupledLite):
+                self._initialize_decoupled_biases()
 
         # Init weights, biases
         initialize_weights(self)
         self.info()
         LOGGER.info("")
+    
+    #新增
+    def _initialize_decoupled_biases(self, cf=None):
+        """Initialize biases for DetectDecoupledLite."""
+        m = self.model[-1]  # DetectDecoupledLite
+
+        # 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."""
@@ -434,7 +534,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, DetectDecoupledLite, 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)

models/yolov5s6_decoupled_light.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

@@ -5,5 +5,6 @@ python /data6/zhangw/yolov5/train.py \
     --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 \
+    --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \
-    --name G0_cub100_yolov5s6_raw_e300
+    --name G8_decoupled_head_light_e300 \
+    --cfg /data6/zhangw/yolov5/models/yolov5s6_decoupled_light.yaml