YOLO检测头自定义改造实战:从P2小目标到注意力机制,附完整步骤

YOLO检测头自定义改造实战:从P2小目标到注意力机制,附完整步骤 1. 检测头到底在做什么先用一句话把它讲明白先说个现象。很多人训练YOLO模型数据标注没问题、训练也能正常跑完、loss也在降但mAP就是卡在某个值上不去。这时候去搜如何提升YOLO精度十有八九会看到改进检测头这个说法。但对零基础的朋友来说光是听检测头三个字就已经开始打退堂鼓了更别提什么Decoupled HeadAnchor-freeDFL这一堆术语。我先用一句话把检测头讲清楚检测头就是网络最后面那一段负责出结果的分支结构它的输入是Backbone和Neck处理好的特征图输出是哪里有个框、框里面是什么类别的预测结果。你可以把整个YOLO网络想象成一条流水线Backbone骨干网络负责看画把原始图像提炼成一张张包含语义信息的特征图。Neck颈部特征融合负责整理信息把不同尺度的特征图融合起来让模型既能看到大物体也能看到小物体。Detect Head检测头负责拍板决策在整理好的特征图上做分类和定位最终输出检测框。很多人一上来就研究怎么改Backbone、怎么调Neck但忽略了检测头。实际上检测头才是决定输出长什么样的组件——你改了Backbone最后输出格式还是那套但只要你动了检测头后面损失函数、NMS后处理、导出部署全都要跟着变。这篇文章就是给零基础的人准备的实战攻略。我会先带你把当前YOLO版本的检测头源码找出来再手把手演示两个改造方向给YOLOv8增加一个P2小目标检测头以及写一个带注意力机制的自定义检测头并替换默认头。全程基于ultralytics官方仓库的代码结构步骤可以直接照着做也会把我踩过的坑一并交代清楚。2. 动手前先定位你的YOLO版本、当前检测头代码和输出shape2.1 版本不一样检测头代码位置完全不同很多人卡在找不到检测头代码这一步其实是因为YOLO的历史版本改动太大了。YOLOv5的检测头在models/yolo.py里YOLOv8和YOLO11的检测头在ultralytics/nn/modules/head.py里YOLOv6是另一个独立的官方仓库结构又是另一套。如果你用的是Ultralytics YOLO系先确认版本import ultralytics print(ultralytics.__version__)我用的版本是8.x代码路径是site-packages/ultralytics/nn/modules/head.py。你只需要在Python环境里执行import ultralytics.nn.modules.head as head print(head.__file__)就能看到这个文件的绝对路径然后用编辑器打开。YOLO11的检测头类名也叫Detect但内部实现有几个细节和v8不一样所以你先确认版本再照葫芦画瓢别直接抄网上v5的改法。2.2 用代码看懂当前的输出结构打开模型看一眼Detect类在哪个位置from ultralytics import YOLO model YOLO(yolov8n.pt) m model.model print(m)最后几行会显示类似(9): Detect( (cv2): ModuleList(...) (cv3): ModuleList(...) (dfl): DFL(...) )cv2是回归分支cv3是分类分支dfl是分布损失相关的解码模块。不同版本里可能还多一个iou分支。再看推理输出的shape这一步对零基础的人极其重要因为它能帮你建立检测头到底输出什么东西的直观概念import torch # 切记切到eval模式因为train模式下Detect的forward逻辑不一样 m.eval() x torch.randn(1, 3, 640, 640) with torch.no_grad(): out m(x) print(out.shape)以COCO 80类的YOLOv8n为例输出shape是[1, 84, 8400]。怎么理解这个数字1是batch size。84 4 80前4个是边界框坐标x1、y1、x2、y2或xywh取决于解码方式后面80个是每个类别的概率。8400是整个输入图像上所有预测位置的总数由三个检测尺度组成80×80 40×40 20×20 8400。如果你的模型加了P2层这个数会变成80×80 40×40 20×20 160×160 34000左右。这个[bs, 4nc, total_predictions]的格式后面所有配套逻辑都围绕它转。你换检测头本质上就是在保证最终输出仍然符合或协调地修改这个格式的前提下改变中间的计算方式。2.3 从yaml到Detect类的加载链路Ultralytics YOLO的模型结构由一个yaml文件描述。以YOLOv8n为例ultralytics/cfg/models/v8/yolov8.yaml的最后几行是head: ... - [-1, 1, Conv, [256, 3, 2]] - [[-1, 24], 1, Concat, [1]] - [-1, 1, C2f, [256, True]] - [8, 6, Conv, [64, 1, 1]] - [-1, 1, nn.Upsample, [None, 2, nearest]] - [[-1, 4], 1, Concat, [1]] - [-1, 1, C2f, [64, True]] - [[14, 27, 33], 1, Detect, [nc]]yaml每一行对应网络的一层其中第三列写的是模块类名比如Conv、C2f、Detect。Ultralytics在构建模型时会去tasks.py的parse_model函数里通过一个全局名字空间的映射找到这些类并实例化。这给了我们一个很重要的结论如果你想用自己写的MyDetect替换默认检测头只需要在yaml的最后一行把Detect改成MyDetect并保证这个类在ultralytics的模块作用域里能被找到。后面第五节我会给出完整注册流程。3. 自定义检测头的设计决策结构改动、shape推演与损失匹配3.1 你想要的替换是哪种程度网上说的替换检测头其实涵盖了三类需求工作量天差地别改动级别示例工作量风险加层增加P2检测头让模型多一个高分辨率检测尺度低只改yaml低损失函数无需改动塞模块在原有Detect的分类或回归分支里插入注意力、轻量化卷积中继承Detect改写forward中需验证shape重写完全放弃默认Detect自己写一套全新的head结构高通常要同步改loss和标签分配高不建议新手直接上手我见过不少初学者一上来就想重写整个检测头结果被loss、标签分配、NMS折磨得体无完肤。我的建议是先在加层和塞模块层面做文章收益高且不容易翻车。如果你确实想重写先搞懂本节后面的shape推演和损失匹配逻辑否则各种报错会让人怀疑人生。3.2 shape推演让每一层前后的张量形状都对得上无论怎么改检测头一条铁律是前向传播的张量形状必须一路对得上。我见过很多人在这一步卡住报错信息五花八门本质都是shape没算清。以YOLOv8默认的Detect为例它接收来自Neck的特征图列表每个特征图的shape类似[B, C, H, W]其中H和W是特征图的宽高C是通道数。检测头内部做的事情可以简化为对每个尺度的特征图做卷积把通道压缩成输出需要的通道数。把特征图从[B, C, H, W]展平成[B, C, H*W]再转置成[B, H*W, C]或保持其他格式。在推理时通过解码把网络输出的原始数值转换成真正的坐标和类别概率。如果你自己设计头那么输出通道数必须和你的预测目标一致。比如一个带objectness分支旧式YOLO的检测头每个位置输出1 4 C_cls个值。一个去掉objectness分支的Anchor-free头YOLOv8默认每个位置输出4 C_cls个值。一个使用DFLDistribution Focal Loss的回归头每个坐标不是直接输出一个值而是输出reg_max个概率分布值典型配置reg_max16。所以回归分支的输出通道是4 * reg_max 64分类分支输出C_cls。下面是一个简单的通道推演示例。假设输入特征图是[1, 128, 160, 160]你想设计一个输出类别数为1的自定义头不加objectness输入: [1, 128, 160, 160] 经过一个3x3卷积out_channels64: [1, 64, 160, 160] 展平: [1, 64, 25600] 转置/重排后: [1, 25600, 64]这个64就是每个预测位置的向量长度。如果你的设计是4个边界框值 60个类别分数那就该是64。如果类别数变了输出通道要跟着调NMS后处理读取的维度也要变。3.3 损失函数与输出定义的匹配关系损失函数必须和头输出语义严格一致这是替换头最容易出问题的地方。我把常见对应关系整理成一张表头输出的成分常用损失说明objectness置信度BCEWithLogitsLossYOLOv5等带obj分支的头需要分类分数BCEWithLogitsLoss多标签分类不是softmax交叉熵直接回归的坐标CIoU L1简单直接常见于重写头DFL分布回归Distribution Focal Loss CIoUYOLOv8/v11默认做法如果你只是像第四节这样加一个P2层还是在用默认Detect的DFL输出那么v8自带损失函数会完全兼容。但如果你写了自定义头后删掉了DFL直接输出坐标值那默认loss里的DFL计算就会报错或导致梯度异常。这个问题我会在第七节展开讲。4. 实战一在YOLOv8里新增一个P2检测头来救小目标4.1 P2层为何能提升小目标检测YOLOv8默认有三个检测层分别叫做P3、P4、P5对应8倍、16倍、32倍下采样。对一张640×640的输入图来说P3层的特征图尺寸是80×80每个格子原本对应原图8×8像素区域。但对于特别小的目标比如高空无人机视角下只有十几个像素的车辆8倍下采样仍然可能把它抹掉不少细节。P2层是4倍下采样特征图尺寸160×160拥有更高的空间分辨率能够保留更多小目标的纹理和边缘信息。加一层P2检测头本质上是让模型多了一只更仔细的眼睛对小目标更敏感。但代价也很明显P2层特征图分辨率高计算量和显存占用都会增加。我实测下来在同样batch size和imgsz下加P2后训练显存大约多占用30%到50%。如果显卡比较紧张可以适当调小batch size。4.2 修改模型yaml让Neck把浅层特征引入检测我们以自定义一个yolov8n-p2.yaml为例。你不用自己从零写整个yaml直接把官方yolov8n.yaml复制一份然后修改head部分。官方v8n的head里有一个重要的特征引出点backbone中一个分辨率较高通常是第4层或第5层输出取决于具体版本的特征图会被引出来经过卷积、上采样、Concat、C2f送到Detect。要加P2整体思路是从backbone里找一层输出分辨率是160×160即原图1/4的特征图。给它接一个1×1卷积做通道对齐再经过上采样或直接连接到检测头。把原来的Detectfrom从[P3, P4, P5]改成[P2, P3, P4, P5]。我给一个相对通用的示例片段注意不同版本的索引号会有差异请对照你本地yaml逐行核对head: # ... 前面保留官网yaml的内容 ... # 假设backbone第4层是160x160的高分辨率特征 - [4, 1, Conv, [64, 1, 1]] # 引出P2需要用到的浅层特征通道对齐到64 - [-1, 1, nn.Upsample, [None, 2, nearest]] - [[-1, 2], 1, Concat, [1]] # 和另一条浅层路径拼接 - [-1, 1, C2f, [64, True]] # 作为P2检测层 # 这里假设上面这个新的C2f层索引是38 - [[14, 38, 33, 27], 1, Detect, [nc]]你需要做两件事一是确认自己版本backbone中真实存在的浅层索引二是记住新增C2f层在整个网络里的索引号最后把它填进Detect的from列表里。如果你不确定索引直接在代码里逐层打印出来核对更稳妥from ultralytics import YOLO model YOLO(yolov8n-custom.yaml) # 用你改好的yaml m model.model for i, (name, layer) in enumerate(m.named_children()): print(i, name, layer.__class__.__name__)打印出来的顺序和yaml中层的顺序一致找到你新增的那层确定它的索引号。4.3 用代码验证新增检测层的shape是否正确改完yaml后先用一行代码看看模型能不能正常构建和输出import torch model YOLO(yolov8n-p2.yaml) m model.model m.eval() x torch.randn(1, 3, 640, 640) with torch.no_grad(): out m(x) print(out.shape) # 如果加了P2最后一位应该是8400 160*160 34000如果你看到最后一位不是34000而是8400说明你虽然新增了分支但Detect层仍然只用了原来的三个检测尺度。这时候检查两处一是Detect的from列表是否包含新增层索引二是Detect实例化的过程中self.nl检测层数是否变成了4。如果from已经正确但输出shape还是不对可以直接看中间特征# 把模型的第9层Detect前一层输入抓住来观察shape m.eval() hooks [] for i, module in enumerate(m.model): if isinstance(module, torch.nn.Module): def hook_fn(module, input, output, idxi): if isinstance(output, (list, tuple)): shapes [o.shape for o in output if isinstance(o, torch.Tensor)] print(flayer {idx}: {shapes}) else: print(flayer {idx}: {output.shape}) hooks.append(module.register_forward_hook(hook_fn)) _ m(x)这段代码会在每个模块执行前打印输入/输出shape具体打印方式可以根据输出格式微调帮你一眼定位是哪一层通道对不上。4.4 P2头的训练和效果验证思路验证结构没问题之后用一个小数据集快速训练对比。我建议先用原有的yolov8n.pt做迁移学习只训练少量轮次观察mAP变化。不要一上来就用大规模数据集全量训练浪费时间。yolo detect train datayour_dataset.yaml modelyolov8n-p2.yaml epochs50 imgsz640 batch8注意这里我把model指向你自己的yaml文件但如果你希望从官方预训练权重开始迁移需要用pretrainedTrue参数或先YOLO(yolov8n-p2.yaml).load(yolov8n.pt)。因为官方权重的检测头没有P2这一层加载时会出现不匹配警告这是正常的ultralytics会忽略多出来的层并随机初始化新增层的参数。关于效果我的经验是P2层对小目标类别AP的提升通常比较明显但不会对所有数据集都生效。要是你的数据集里绝大多数目标本身就不小那加了P2反而可能因为引入更多背景噪声导致轻微掉点。判断标准很简单先画出GT框尺寸分布如果大量目标的长边不到原图尺寸的10%P2值得加。5. 实战二写一个带注意力机制的自定义检测头并替换默认头5.1 先想清楚你到底需不需要换整个头我见过一些同学看了几篇论文就照着注意力机制检测头的图去魔改。这里先说句逆耳的对绝大多数常规目标检测任务来说替换检测头带来的收益不如调数据增强、调损失权重、加训练时长来得明显。如果你只是想涨点先把P2、数据增强、超参搜索试完再说。但如果你恰好有检测头输出结构要适应特殊任务的需求——比如要输出旋转框、要输出计数密度图、要融合多个传感器特征——那就真的得换。我以一个给默认检测头加一个轻量通道注意力SE模块的例子演示完整替换流程。这个例子的定位是给零基础一个安全入门的路径因为它继承官方Detect类不改变既有输出格式所以损失函数完全不用动。5.2 写一个继承官方Detect的SEAttnDetect思路很简单在官方Detect前处理特征图时先过一个SE块对通道做注意力重标定然后再走默认的检测分支。import torch import torch.nn as nn from ultralytics.nn.modules import Detect class SEBlock(nn.Module): 标准Squeeze-and-Excitation模块先用全局平均池化压缩空间信息再通过两个1x1卷积建模通道间依赖。 def __init__(self, channels, reduction4): super().__init__() self.pool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Conv2d(channels, channels // reduction, 1, biasFalse), nn.ReLU(inplaceTrue), nn.Conv2d(channels // reduction, channels, 1, biasFalse), nn.Sigmoid(), ) def forward(self, x): # x: [B, C, H, W] weight self.pool(x) # [B, C, 1, 1] weight self.fc(weight) # [B, C, 1, 1] return x * weight class SEAttnDetect(Detect): def __init__(self, nc80, ch()): super().__init__(nc, ch) # 为每个检测尺度配一个SE模块ch里存的是各个尺度特征图的通道数 self.se nn.ModuleList([SEBlock(c) for c in ch]) def forward(self, x): # x是来自Neck的特征图列表 x_attn [] for i, feat in enumerate(x): x_attn.append(self.se[i](feat)) # 继续走官方Detect的forward包括训练/推理分支 return super().forward(x_attn)你可能会问为什么不在cv2、cv3分支之后各挂一个SE而是在整个特征图进Detect之前挂原因是后者改动更小也能体现注意力对多尺度特征的微调对刚入门的人更友好。真正工业级的优化往往会把SE放进分支内部但那是后话。5.3 注册自定义头让yaml能识别SEAttnDetectUltralytics构建模型的parse_model函数会从一个全局名字空间里查找类名。最直接的办法就是把上面的代码追加到ultralytics/nn/modules/head.py文件末尾然后在ultralytics/nn/modules/__init__.py中加入导出from .head import SEAttnDetect如果你不想改原始安装包也可以在启动脚本里手动把模块注入到ultralytics的模块命名空间。我推荐前者干净直接。然后创建一个自定义yaml可以基于yolov8n-p2.yaml再改把最后一行从- [[14, 38, 33, 27], 1, Detect, [nc]]改成- [[14, 38, 33, 27], 1, SEAttnDetect, [nc]]这里有个细节SEAttnDetect的初始化签名是(nc, ch)parse_model会自动把nc和检测层传入的ch列表传给你所以你的类定义里这两个参数不能少。然后验证输出shape是否和官方一致model YOLO(yolov8n-se.yaml) m model.model print(m.model[-1]) # 应该能看到SEAttnDetect如果打印出来最后一层确实是SEAttnDetect那恭喜你替换检测头这件事已经跑通了大半。5.4 使用自定义头时最容易忽略的forward模式问题如果你在训练时报错或者eval时输出shape和官方不一样先不要怀疑代码写错极大概率是Detect父类的forward内部对self.training做了分支处理。继承后要确保你的自定义类没有在forward里额外改变训练/推理的状态。我们的示例里super().forward(x_attn)会自己处理所以没问题。另一个常见的坑是很多人在自定义头里忘掉self.stride的更新。Detect类在初始化时会通过self.stride保存每个检测层对应的下采样倍数如[8, 16, 32]训练时Model会在model.train()前自动调用_initialize_detector等逻辑来设置stride。如果你重写了__init__但没有完整调用父类的初始化逻辑stride可能就是空的推理时输出的坐标会异常放大到离谱的数值。我们的示例直接调用super().__init__所以避开了这个问题。6. 检测头换完之后这三个配套模块必须同步调整很多人把检测头替换完看着训练能跑就以为大功告成结果部署一测坐标全乱或者mAP反降。这类问题的根源通常是三个配套模块没跟上损失函数、标签分配、后处理与部署格式。6.1 损失函数你对齐了吗YOLOv8默认的损失函数v8DetectionLoss是按照官方Detect的输出结构设计的。如果你换的头仍然输出DFL分布分类分数那它天然兼容。但如果你删掉了objectness或DFL就必须同步修改ultralytics/utils/loss.py中的对应部分。我建议的排查顺序是先跑一个极小的数据集观察loss曲线是否正常下降。如果loss出现NaN或者训练初期就异常大八成是损失函数和输出结构不匹配。举例来说你的自定义头没有objectness输出但Loss里还留着self.bce(obj)的计算就会报维度不匹配的错误。6.2 标签分配Anchor-free和Anchor-based不能混用YOLOv8默认是Anchor-free结构训练时使用TaskAlignedAssigner做标签分配它根据分类分数和IoU的加权结果决定正样本。如果你换成了一个Anchor-based的检测头类似YOLOv5那标签分配逻辑必须换成基于Anchor的分配器否则模型根本学不到正确的正样本匹配。更隐蔽的问题是加了P2层之后特征图尺度变多标签分配时的scale权重不同TaskAlignedAssigner的行为也可能需要微调。我见过有人在加了P2后小目标AP不升反降最后发现不是模型结构问题而是新尺度的正样本分配和老尺度产生冲突导致某些小目标变成负样本被抑制了。6.3 后处理与导出部署ONNX、TensorRT、FPGA各有脾气在训练脚本里NMS后处理是ultralytics自动完成的但导出后就要自己处理输出。以ONNX为例你导出的模型输出shape仍然是[bs, 4nc, total_predictions]你需要自己写解码NMS逻辑。如果你的自定义头引入了新算子比如SE模块里常见的ReduceMean和卷积算子标准ONNX和TensorRT一般都能支持但如果你用了复杂的自定义autograd.Function导出基本必炸。这里特别提一句FPGA等端侧场景很多FPGA加速器不支持动态shape导出时需要固定输入分辨率。另外DFL层的解码过程是纯后处理建议从模型里拆出来放到上位机或CPU端做而不是留在加速器里硬算否则latency会很难看。我在板端部署时习惯先python export.py导出ONNX再用onnxsim简化最后手动验证每个输出节点的shape和语义。7. 替换过程中最容易踩的三个坑与排查链路7.1 坑一yaml改了但实际跑的还是旧模型这个坑我踩过不止一次。你改了yaml但是训练时传入的modelyolov8n.pt它会直接加载权重里的模型结构完全无视新yaml。技术上讲这是两套构建路径# 错误示范加载权重文件结构是权重里的旧结构 model YOLO(yolov8n.pt) # 正确示范用yaml构建模型结构再手动加载权重 model YOLO(your-custom.yaml).load(yolov8n.pt)如果你发现改了yaml但输出shape没变化第一件事先检查你到底传给YOLO()的是yaml还是pt。7.2 坑二Loss突然变成NaN自定义头训练后Loss变成NaN90%的情况出在数值上。排查链路我一般按这个顺序走检查输入标签是否存在空的ground truth框。可以用一行代码检查torch.isnan(targets).any()如果有NaN问题在数据标注不在模型。检查自定义模块里有没有exp、log这类容易溢出的算子。尤其是自己写的DFL解码逻辑数值范围控制不好直接炸。检查学习率。自定义头相当于引入了新的随机初始化参数和学习率不匹配时前期容易把loss冲飞。建议把初始学习率降到默认值的1/5试跑20个step看loss是否稳定。在forward里手动打印每一层输出的数值范围定位inf/NaN是从哪一层开始的。torch.autograd.set_detect_anomaly(True)这行代码可以帮你把异常梯度出现的位置打出来非常有用。7.3 坑三小目标AP不升反降加了P2之后小目标AP反而下降看起来很反直觉但其实有很多合理原因浅层特征噪声大。160×160的特征图包含丰富的空间细节但也包含大量背景纹理。如果模型容量不够容易把背景错检成目标。正样本分配失衡。新增的P2层会产生大量候选位置导致匹配到的正样本中大部分来自浅层其他尺度学不到东西。数据增强不匹配。小目标本身在Mosaic等增强下容易被裁剪掉P2层再强也无法无中生有。我建议用Ultralytics训练时自带的PR曲线和混淆矩阵去判断如果P2带来的不是目标被检出而是误检增多那就把confidence阈值调高或者给P2层单独加一点输出的偏置惩罚如果小目标类别本身的召回率提升但精度下降可以考虑用focal loss调整难易样本权重。7.4 我的建议调试顺序如果你现在正对着一个不收敛或shape报错的模型不要东一榔头西一棒子地试。我自己的流程是用CPU 极小数据集几十张图 1个epoch跑通前向和反向确认没有shape和数值错误。固定随机种子用同样数据分别训练官方模型和改头模型各跑10个epoch看loss差异。结构没问题后再上GPU、加大数据、调超参。这一套流程看着保守但实际上最省时间。那些改了头就一下提升好几点的传奇案例背后大概率是作者已经理解了结构、损失、标签分配之间的联动关系。零基础的朋友按这个顺序来至少不会在第一步就被报错劝退。