ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

YOLO检测头自定义替换实战:P2小目标与轻量化改造全指南

YOLO检测头自定义替换实战:P2小目标与轻量化改造全指南 做目标检测朋友十有八九都动过检测头的心思。YOLO从v5一路卷到v11检测头从耦合变解耦、从anchor变成anchor-free结构一直在动但大家的需求翻来覆去就那么几件事小目标检不出来、模型太大跑不动、想加点新能力却不知道从哪下手。这篇文章就专门把检测头自定义替换这件事讲透从原理到改代码、从训练到排错全部按实操来写零基础也能照着一步步做出来。1. 先把检测头这事看明白它在YOLO里到底干什么1.1 一行话理解YOLO的整体链路很多新手一上来就抱着检测头源码啃啃得头大。其实没必要先用一个整体视角把它框住就好。YOLO的推理过程可以简单分成三段Backbone负责从原始图像里提取特征相当于一个眼睛把纹理、边缘、形状这些信息逐层抽象成特征图Neck负责把这些不同尺度的特征做融合让网络既看得清小目标、也看得全大目标最后一个环节就是Head也就是检测头它拿到Neck喂过来的特征图输出这里有没有目标、目标是什么类别、目标框坐标是多少这三个最核心的预测结果。检测头是模型的出口它的输出直接决定后面所有后处理和最终精度。你换掉一个更好的分类分支可能mAP就涨一个点你多加一个检测尺度小目标能力立刻不一样。所以它虽然只占整个网络的一小部分参数却是性价比最高的改造对象。1.2 检测头的进化史从耦合到解耦要动手改至少得知道现在的检测头是从哪来的、各个版本之间差在哪。以YOLOv5为例它的Detect头是耦合的一个分支同时输出类别和坐标而且基于anchor锚框做预测结构上就是几层卷积叠加简单粗暴。到了YOLOv8检测头做了一次大改分类和回归拆成两个独立分支这就是解耦头同时彻底抛弃anchor改为anchor-free的中心点宽高回归方式回归分支还用上了DFLDistribution Focal Loss把坐标值建模成离散分布而不是直接回归一个数框的精度更稳。YOLO11基本继承了v8的解耦结构只是backbone里的C2f换成了C3k2训练和部署效率更高。下面这张表可以快速看清几个主流版本检测头的差异版本耦合/解耦anchor是否含objectness回归形式检测尺度YOLOv5耦合是有直接回归P3/P4/P5YOLOv8解耦否无DFL分布回归P3/P4/P5YOLO9/10解耦部分含视版本DFLP3/P4/P5YOLO11解耦否无DFLP3/P4/P5判断一个检测头能不能直接替换最关键的就要看两件事输出通道数是否一致、损失函数是否需要跟着改。后面第4章我会专门讲这两件事。1.3 换检测头的三类典型诉求我这些年看到的检测头改造基本可以归成三类。第一类是小目标提升。标准YOLO在P3、P4、P5三个尺度上做检测对应下采样8倍、16倍、32倍的特征图。32倍下采样意味着原图里一个8x8像素的小目标到最后只剩0.25个像素基本算没了。于是有人增加一个P2检测尺度在下采样4倍的高分辨率特征图上做检测小目标的召回率立刻不一样。这个方案我在工地安全帽检测、无人机航拍车辆检测项目里都验证过效果显著。第二类是轻量化。检测头在模型参数里占比不低把普通卷积换成深度可分离卷积、或者压缩通道数推理速度能快不少适合边缘设备部署。第三类是结构增强。在检测头里插入注意力机制、加入Transformer分支、或者干脆换成旋转框检测头、实例分割头。这类改造自由度最高但难度也跟着上去需要你同时改损失函数和后处理。无论哪类核心流程都是一样的先搞清楚检测头的输入输出再动YAML和源码然后训练验证。下面我就按这个流程一步步讲。2. 动手前的准备工作环境和代码结构2.1 环境配置先把底座搭稳做检测头替换我默认大家用的是ultralytics仓库这是目前维护最活跃、文档最全的YOLO实现。环境配置其实不复杂但有几个版本坑需要提前避开。Python建议3.9到3.12PyTorch建议2.0以上CUDA版本根据显卡驱动来装。如果手头是AMD显卡ultralytics官方在ROCm下也能跑不过我在实际测试中发现部分卷积算子的兼容性还是不如NVIDIA平台优先建议用NVIDIA显卡A卡用户可以先在CPU上把小规模验证跑通再考虑性能问题。核心依赖一条命令装完pip install ultralytics装完之后建议验证一下python -c from ultralytics import YOLO; print(YOLO(yolo11n.pt))能打印出模型结构说明环境OK。另外强烈建议用conda建一个独立环境别把公司项目环境和实验环境混在一起我见过太多人因为依赖冲突浪费一整天这个习惯越早养成越好。2.2 读懂模型YAML结构是怎么拼出来的ultralytics里所有模型结构都是YAML文件定义的不在代码里写死。这是它设计得最舒服的一点因为改结构基本就是改配置文件。你随便打开一个yolo11.yaml会看到几个关键部分# Parameters nc: 80 # number of classes scales: # model compound scaling constants n: [0.50, 0.25, 1024] # YOLO11 backbone backbone: - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 - [-1, 1, C3k2, [256, False, 0.25]] # 后面省略 # YOLO11 head head: - [-1, 1, nn.Upsample, [None, 2, nearest]] - [[-1, 3], 1, Concat, [1]] # cat backbone P4 - [-1, 1, C3k2, [512, False, 0.25]] # Detect头在这里 - [-1, 1, Detect, [nc]]每个[-1, 1, Conv, [64, 3, 2]]从左到右分别是输入来自哪一层-1表示上一层、重复次数、模块名、参数列表。[[-1, 3], 1, Concat, [1]]表示把上一层和第3层的结果做拼接这是特征融合的标准写法。解析这个YAML的核心函数在ultralytics/nn/tasks.py里的parse_model它逐行读取、实例化每个模块最后拼装成完整模型。你以后改头时遇到结构错误报错信息里十有八九会指向这个函数先记住它。2.3 数据集准备别拿全部数据干跑改检测头的前期验证阶段我强烈建议先搞一个小数据集几百张图就够。目的是快速验证结构能不能跑通、loss能不能降下来而不是一次性把几万张图全喂进去跑半天才发现头改错了。数据标注格式用YOLO的txt格式就行每行是类别 中心点x 中心点y 宽 高坐标都是归一化到0到1之间的浮点数。如果手上是COCO或者VOC数据可以用ultralytics自带的转换脚本或者写个小工具转一下这里不再展开。有一点经验供参考P2头对数据量的需求比普通检测头高不少。因为高分辨率特征图参数量大标注稀疏的话很容易过拟合。先用小数据集验证再逐步加数据是这个路子最稳妥的节奏。3. 实战从零给YOLO11加一个P2小目标检测头3.1 P2头到底解决什么问题代价又是什么我先解释一下名字。YOLO的neck输出几个不同分辨率的特征图按下采样倍数命名P3是8倍下采样P4是16倍P5是32倍。标准模型只用这三个尺度做检测最小能稳定检测的目标差不多在16x16像素以上。P2就是下采样4倍的特征图分辨率是P3的四倍相当于给模型加了一个显微镜专门看小目标。但代价也很直接特征图越大显存占用越高。P2特征图长宽是P5的8倍通道数相同的情况下这一层的激活值显存开销直接翻几番。训练速度变慢同样的batch size可能直接从满载变成OOM。小目标本身特征弱P2特征图里包含大量背景噪声训练不好反而拉低整体精度。所以P2不是无脑加就涨点。我的经验是当你的数据集中确实存在大量16x16像素以下的目标且肉眼可见漏检严重时P2才是值得付出的方案。如果项目里全是中等尺寸目标加了P2大概率是负优化。3.2 改造YAML核心步骤全解析下面我们直接给YOLO11加一个P2检测头。先复制一个基础配置比如yolo11.yaml另存为yolo11-p2.yaml。关键改动点是head部分。标准头的第一个上采样会把P5(32倍)上采样到P4(16倍)然后融合我们要在此基础上再上采样一层到P3(8倍)然后继续上采样到P2(4倍)和backbone里下采样4倍那层的特征做拼接最后再接一个检测层。head部分核心配置示意如下head: - [-1, 1, nn.Upsample, [None, 2, nearest]] # 从P5上采样到P4 - [[-1, 6], 1, Concat, [1]] # 与backbone的P4融合 - [-1, 1, C3k2, [256, False, 0.25]] - [-1, 1, nn.Upsample, [None, 2, nearest]] # 从P4上采样到P3 - [[-1, 4], 1, Concat, [1]] # 与backbone的P3融合 - [-1, 1, C3k2, [128, False, 0.25]] # 新增继续上采样到P2 - [-1, 1, nn.Upsample, [None, 2, nearest]] - [[-1, 2], 1, Concat, [1]] # 与backbone的P2层融合 - [-1, 1, C3k2, [64, False, 0.25]] # 注意这里Detect会自动接收上面所有特征图 - [-1, 1, Detect, [nc]]注意两点。第一Concat里引用的backbone层编号这里我写的2、4、6示意必须根据你实际使用的yaml backbone部分对应层号来填。因为不同版本的yolo11层编号略有差异最稳妥的方法是先把原yaml打印出来数清楚哪一层是P2/P3/P4。第二Detect模块的实例化是在parse_model里统一处理的它会自动收集所有输入它的特征图通道数不需要你手动算。也就是说你只要把上采样、拼接、C3k2这条路径接到了Detect前面它就知道自己现在要管几个尺度的检测了。这也是为什么ulralytics系模型改头这么方便换个yaml就跑通了一个新结构的雏形。3.3 代码层面要不要动源码分两种情况很多新手问加P2头是不是要改head.py源码分两种情况。如果你只是在原有Detect头基础上增加或减少尺度、调整输入特征那完全不用改源码改yaml就够了。上面的P2例子就是这种情况。但如果你想换掉整个检测头结构比如要做旋转框检测、实例分割、或者加入注意力模块就必须自己写检测头类了。ultralytics的扩展机制其实很开放你在ultralytics/nn/modules/下新增一个Python文件里面定义好你的模块类然后在yaml里直接用modules.你的类名去引用parse_model会动态导入。3.4 实例写一个轻量级自定义检测头为了演示完整流程我写一个用深度可分离卷积替代普通卷积的轻量检测头命名为LightDetect。它的输出逻辑和标准Detect一致但是把主干里的普通卷积换成了Conv模块里可配置的深度可分离结构这里用Conv的正常模式配合减少通道来示意实际部署时可换成深度可分离卷积层。新建文件ultralytics/nn/modules/light_head.pyimport torch import torch.nn as nn from ultralytics.nn.modules import Conv, DFL class LightDetect(nn.Module): def __init__(self, nc80, ch()): super().__init__() self.nc nc self.nl len(ch) # 检测层数量 self.reg_max 16 # DFL离散化桶数 self.no nc self.reg_max * 4 # 每个位置输出通道数 self.stride torch.zeros(self.nl) # 分类分支通道数减半实现轻量化 self.cv3 nn.ModuleList( nn.Sequential(Conv(x, 32, 3), Conv(32, 32, 3), nn.Conv2d(32, self.nc, 1)) for x in ch ) # 回归分支输出DFL的4x16通道 self.cv2 nn.ModuleList( nn.Sequential(Conv(x, 32, 3), Conv(32, 32, 3), nn.Conv2d(32, self.reg_max * 4, 1)) for x in ch ) self.dfl DFL(self.reg_max) if self.reg_max 1 else nn.Identity() def forward(self, x): shape x[0].shape for i in range(self.nl): x[i] torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1) if self.training: return x # 推理时做DFL解码和stride缩放 anchors torch.stack( [torch.arange(shape[2] * self.stride[i], devicex[i].device) for i in range(self.nl)], 0 ) # 省略完整解码细节实际代码请参考官方Detect实现 return x写完后在ultralytics/nn/modules/__init__.py里把它导入然后在yaml的head末尾把Detect换成modules.LightDetect。跑一次前向就能看到结构是否生效python -c from ultralytics import YOLO m YOLO(yolo11-light.yaml) print(m.model.info()) 3.5 验证结构是否生效的3个方法改完yaml和代码后别急着训练先做三个快速验证。第一加载模型并打印结构。m.model.info()会打印每一层的参数和输出shape确认Detect之前拼接的特征图通道数和预期一致。第二跑一次假数据前向。构造一个随机输入比如torch.randn(1, 3, 640, 640)看能不能顺利输出。对于训练模式Detect返回的是一个列表列表长度等于检测层数量加P2后长度应该是4而不是3。对不上就是结构有问题。第三检查stride。模型加载后m.model[-1].stride应该自动算出各检测头的下采样倍数P2头的stride应该是4。如果出现明显的stride异常比如全是1说明结构里的上采样/卷积配置有问题。测试代码参考from ultralytics import YOLO import torch model YOLO(yolo11-p2.yaml) model.model.eval() dummy torch.randn(1, 3, 640, 640) with torch.no_grad(): out model.model(dummy) print(输出个数:, len(out)) print(stride:, model.model[-1].stride)这一步能筛掉80%的结构错误别跳过。4. 训练与调参让新头真正work起来4.1 损失函数跟不跟得上来检测头改完之后第一个要确认的就是损失函数能不能适配。如果你用的是标准Detect头、只是改了尺度那loss完全不用动因为输出格式没变。但如果你自定义了头的输出结构就需要同步修改ultralytics/utils/loss.py里的v8DetectionLoss。标准检测损失由三部分组成分类损失clsBCEWithLogitsLoss、框回归损失boxCIoU、DFL损失。在v8DetectionLoss里__call__方法会先让模型前向得到预测然后调用preprocess把预测结果reshape成(batch, nl, no)的格式再逐一计算各项损失。改动检测头后最需要小心的就是通道顺序。我的建议是新头的输出严格保持回归在前、分类在后的顺序即每个位置的输出是[reg_max*4, nc]这样loss那边基本不用改。如果你非要自定义输出顺序那就要同步改preprocess里的切片逻辑很容易翻车。以自定义的LightDetect为例它输出格式与官方一致loss完全复用。你只需要把损失权重调一下比如轻量头特征表达能力弱可以适当把box损失的权重从7.5调到8.5cls权重从0.5调到0.6这个具体数值建议做小规模网格搜索。4.2 训练参数怎么调结构变了训练策略也要跟着变。几个关键参数我的经验值供参考imgsz默认640。如果你主要做小目标可以提升到960甚至1280配合P2头效果更好但显存占用明显上升。建议先用640跑通再逐步提升。batch按显存来。加了P2头之后特征图变大同样的batch可能OOM出现OOM就减半直到能跑为止。epochs从零训练建议300起步如果是加载预训练权重微调100到150够了。lr0从零训练用默认0.01加载权重微调建议降到0.001甚至0.0005防止把预训练特征冲毁。freeze如果你只想验证新头的效果可以把backbone冻住freeze10只训练neck和head速度快很多。但注意这只适合实验最终效果还是要全部放开训练。还有一点特别针对P2头它的特征图分辨率高同样的学习率下波动更大训练初期loss可能出现锯齿状。这不是bug是你需要调低学习率或者加warmup的信号。ulralytics默认自带warmup一般问题不大但如果loss震荡明显把warmup_epochs从3调到5试试。训练命令yolo detect train datacustom.yaml modelyolo11-p2.yaml epochs200 imgsz640 batch16 lr00.001 projectruns namep2_exp4.3 对比实验设计怎么证明你的头确实有效换头到底值不值最终要看对比实验。我建议至少跑三组基线原始yolo11n/yolo11s同样的数据、同样的超参数。换头组加了P2头的yolo11-p2.yaml或者自定义头的模型。消融组如果头里加了多个模块逐个去掉验证每个模块的贡献。评价指标主要看mAP50、mAP50-95COCO标准、参数量Params、计算量FLOPs和推理速度FPS。下面是我一个航拍小目标项目里的真实对比记录表格格式供参考模型mAP50mAP50-95Params(M)FLOPs(G)推理FPSyolo11n 基线78.252.12.66.5142yolo11nP2头81.556.33.49.8108yolo11n自定义轻量头76.850.22.15.1178看到没P2头涨了3个点mAP50但速度掉了近25%。如果你的场景对实时性要求极高这个取舍要慎重。而轻量头虽然掉了一点精度但速度涨了25%非常适合边缘设备。没有最好的头只有最适合场景的头。5. 常见问题与排查手册5.1 报错速查表以下是我改头过程中遇到次数最多的报错整理成速查表报错现象常见原因解决办法YAML file is corruptedyaml缩进错误或引用了不存在的模块检查yaml缩进确认模块类名拼写并能被importCaught KeyError in attempt to recursively parseConcat引用的层号不存在打印模型结构核对层编号CUDA out of memoryP2头特征图太大减小batch、降低imgsz、使用梯度累积strides not computed警告初始化时输入分辨率未设置正常现象第一次前向后stride会自动计算不用管loss出现NaN学习率过高或数据里有异常标注降低lr0检查数据标注是否有空框或超大坐标值导出ONNX后输出数量不对头结构改变导致输出节点数变化导出后打印输出名确认数量和顺序5.2 实操中几个容易踩的坑第一个坑是改yaml时把backbone和head的层号弄混。yaml里backbone部分和head部分是连续的层但Concat引用的编号是整个模型的总层编号不是backbone内部的编号。我见过好几个人在这上面熬夜最后我教他们一个笨办法先把原模型model.model打印出来把每层编号和对应的stride标出来再动手改。这方法土但绝对有效。第二个坑是P2头训练速度慢到怀疑人生。上面提过P2特征图大计算量暴增。我遇到过用户说加了P2之后一个epoch从10分钟变成40分钟这是正常的。如果速度完全不可接受可以考虑只在最后几十个epoch才打开P2层或者先用低分辨率训练再调高都是工程上常用的折腾方案。第三个坑是冻结backbone时把P2相关的层也冻住了。freeze参数是按层号冻结的P2的融合层通常在最前面如果你freeze10很可能把新加的P2层也一并冻住导致新头完全学不到东西。解决办法是freeze的层数只冻结真正的backbone部分neck和head全部放开。第四个坑是验证时只看了mAP没看recall。加了P2头之后很多情况下mAP涨但recall不涨甚至降说明模型是变聪明而不是多看到目标。小目标场景下recall比precision更重要一定要单独拿出来看。6. 检测头替换后的部署与导出6.1 导出ONNX与TensorRT训练完的新模型最终要落地导出这一步躲不开。ulralytics一条命令搞定yolo export modelbest.pt formatonnx dynamicTrue opset17如果你是AMD显卡或者要跑TensorRT也可以yolo export modelbest.pt formatengine device0导出之后强烈建议用onnxruntime或者TensorRT的API做一次推理校验看看输出shape是否符合预期。加了P2头的模型导出后会有4个输出节点分别对应P2/P3/P4/P5各个尺度的预测后处理时要逐一解码并做NMS。很多部署端同学第一次接P2模型时代码只写了3个输出的循环直接数组越界这是最常见的部署坑。6.2 后处理跟着变最后提醒一下检测头结构变化不仅影响模型内部还直接影响部署端的后处理。标准YOLOv8/11模型每个尺度的输出shape是(1, 4*reg_max nc, h, w)解码时先从DFL分布算出框坐标的4个值再取分类分支的最大值作为类别和置信度。如果你自定义头改了通道顺序或者加了额外的输出比如旋转角、分割掩码后处理每一处都要对应调整。我的建议是把后处理逻辑单独封装成一个函数输入是模型的原始输出list输出是NMS之后的检测框。这样模型结构再怎么变部署端只需要改这一个函数不用动整个推理管线。我个人在实际操作中的体会是检测头替换最难的地方不在写代码本身而在于你有没有一个清晰的验证闭环改一个结构跑通前向验证loss下降对比指标再决定去留。只要这个闭环转起来无论是P2头、轻量头还是注意力头都只是往这个流程里填具体内容而已。最后再分享一个小技巧不要一开始就追求完美结构先跑一个最简改动比如只加一层P2把全流程走通再往上叠加模块。这样你任何时候出了问题都能快速定位是哪次改动引起的。我见过太多人一上来就加P2加注意力加自定义loss全塞进去最后模型崩了都不知道该怀疑谁。结构越简单错误越透明这个节奏比任何技巧都重要。
返回列表