From 0bf2989264a99e33ae465b8e2e87eafe4dd41826 Mon Sep 17 00:00:00 2001 From: Zhang Wei <3160665918@qq.com> Date: Fri, 20 Mar 2026 13:38:48 +0800 Subject: [PATCH] cls_head_enhance_and_ECA --- models/yolo.py | 130 ++++++++++++++++++++++++++- models/yolov5s6_cls_enhance_ECA.yaml | 62 +++++++++++++ train_my_gra.sh | 4 +- 3 files changed, 191 insertions(+), 5 deletions(-) create mode 100644 models/yolov5s6_cls_enhance_ECA.yaml diff --git a/models/yolo.py b/models/yolo.py index db7ac21..14183ce 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -138,6 +138,24 @@ class ClsEnhanceBlock(nn.Module): def forward(self, x): return x + self.cv2(self.cv1(x)) +# 新增 +class ECABlock(nn.Module): + """Efficient Channel Attention for G12.""" + 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 DetectDecoupledLite(nn.Module): stride = None @@ -290,6 +308,87 @@ class DetectDecoupledLiteClsEnhance(nn.Module): anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape) return grid, anchor_grid +#新增 +class DetectDecoupledLiteClsEnhanceECA(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 enhance + self.cls_enhance = nn.ModuleList(ClsEnhanceBlock(head_c, expand=0.5) for _ in ch) + + # G12 ECA attention + self.cls_attn = nn.ModuleList(ECABlock(head_c, k_size=3) 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) + 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): """YOLOv5 Segment head for segmentation models, extending Detect with mask and prototype layers.""" @@ -373,7 +472,8 @@ class BaseModel(nn.Module): m = self.model[-1] # Detect() # if isinstance(m, (Detect, Segment)): # if isinstance(m, (Detect, DetectDecoupledLite, Segment)): - if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)): + # if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)): + if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, DetectDecoupledLiteClsEnhanceECA, Segment)): m.stride = fn(m.stride) m.grid = list(map(fn, m.grid)) if isinstance(m.anchor_grid, list): @@ -412,7 +512,8 @@ class DetectionModel(BaseModel): m = self.model[-1] # Detect() # if isinstance(m, (Detect, Segment)): # if isinstance(m, (Detect, DetectDecoupledLite, Segment)): - if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)): + # if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment)): + if isinstance(m, (Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, DetectDecoupledLiteClsEnhanceECA, Segment)): def _forward(x): """Passes the input 'x' through the model and returns the processed output.""" @@ -431,6 +532,8 @@ class DetectionModel(BaseModel): self._initialize_decoupled_biases() elif isinstance(m, DetectDecoupledLiteClsEnhance): self._initialize_decoupled_cls_enhance_biases() + elif isinstance(m, DetectDecoupledLiteClsEnhanceECA): + self._initialize_decoupled_cls_enhance_eca_biases() # Init weights, biases initialize_weights(self) @@ -477,6 +580,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_eca_biases(self, cf=None): + """Initialize biases for DetectDecoupledLiteClsEnhanceECA.""" + m = self.model[-1] + + # objectness bias + 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) + + # classification bias + 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: @@ -650,7 +773,8 @@ def parse_model(d, ch): # TODO: channel, gw, gd # elif m in {Detect, Segment}: # elif m in {Detect, DetectDecoupledLite, Segment}: - elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment}: + # elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, Segment}: + elif m in {Detect, DetectDecoupledLite, DetectDecoupledLiteClsEnhance, DetectDecoupledLiteClsEnhanceECA, 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) diff --git a/models/yolov5s6_cls_enhance_ECA.yaml b/models/yolov5s6_cls_enhance_ECA.yaml new file mode 100644 index 0000000..fbfd98a --- /dev/null +++ b/models/yolov5s6_cls_enhance_ECA.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, DetectDecoupledLiteClsEnhanceECA, [nc, anchors]], # Detect(P3, P4, P5, P6) + ] \ No newline at end of file diff --git a/train_my_gra.sh b/train_my_gra.sh index 0d6d1ee..402a88a 100644 --- a/train_my_gra.sh +++ b/train_my_gra.sh @@ -6,5 +6,5 @@ python /data6/zhangw/yolov5/train.py \ --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \ --weights /data6/zhangw/yolov5/yolov5s6.pt \ --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \ - --name G11_cls_head_enhance_Lightweight_e300 \ - --cfg /data6/zhangw/yolov5/models/yolov5s6_cls_enhance.yaml \ No newline at end of file + --name G12_cls_head_enhance_ECA_e300 \ + --cfg /data6/zhangw/yolov5/models/yolov5s6_cls_enhance_ECA.yaml \ No newline at end of file