这次我们继续 RK3588 实战系列。第四集要处理的问题非常具体YOLOv8 默认的输出头是单头输出如何在源码层面把它改成双头输出并且顺利落到 RK3588 的 RKNN 环境里完成转换和板端验证。这不是一个概念演示而是一条可以照着操作的完整链路。默认情况下YOLOv8 的 Detect 头会把边界框回归和类别分类拼接在一个张量里输出最终 shape 可以理解为1, 8400, 4nc。在 PyTorch 里用起来很自然但到了 RK3588 的 NPU 部署场景很多人会碰到两个痛点一是量化精度不容易调二是后处理要在混合张量里反复做切片。改成双头输出本质就是把一个混合张量拆成两个独立张量一个专门输出边界框一个专门输出分类分数。这样的结构在板端解析时更清晰也为后续扩展多任务头留好了位置。本文会把整个改造过程拆成四步来展开拆解 YOLOv8 输出层的源码结构、在源码层把单头改成双头、导出 ONNX 并确认输出节点、在 RKNN 工具链里完成转换和板端验证。整个过程不需要重新训练模型只是对推理图和后处理方式做调整适合正在 RK3588 上做边缘检测部署的读者。1. 核心能力速览能力项说明项目类型RK3588 边缘端目标检测部署改造核心操作拆解 YOLOv8 Detect 输出层将单头输出改造为双头输出目标平台RK3588 开发板配合 RKNN-Toolkit2 工具链前置条件YOLOv8 源码、PyTorch 环境、ONNX、rknn-toolkit2主要功能同时输出边界框回归张量与类别分数张量便于分离后处理模型训练不需要重新训练只调整推理结构推理设备RK3588 内置 NPU是否支持接口板端可通过 RKNN Python API 或 C API 封装推理服务是否支持批量任务板端可以逐帧或分批推理多路摄像头需要自行设计任务队列改造难度中等要求能看懂 PyTorch 网络定义和基本张量操作整个改造的核心不在 RKNN 转换脚本里而在 YOLOv8 的 Detect 头代码里。只要把输出张量的切分逻辑理清楚后续 ONNX 导出和 RKNN 转换都很顺。2. 单头输出与双头输出的本质区别2.1 YOLOv8 默认的单头输出结构YOLOv8 的检测头分为三个尺度分别对应 P3、P4、P5 特征图。每个尺度都会经过两组卷积分支cv2负责边界框回归cv3负责类别分类。在默认的Detect.forward里这两组卷积的输出会按通道维度拼接再把三个尺度展平后的结果合并成一个大张量输出。# 简化后的 YOLOv8 Detect forward 核心逻辑不同 Ultralytics 版本行号会有差异 for i in range(self.nl): # cv2: 回归分支cv3: 分类分支 x[i] torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1) # 把三个尺度的特征图展平并拼接 x_cat torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2) # 最终输出 shape 为 (1, 8400, 4 * reg_max nc) 或 (1, 8400, 4 nc) return x_cat.transpose(1, 2)这里有一点需要注意不同版本的 Ultralytics 对 DFL 解码的位置处理并不一致。有的版本会在Detect.forward里先完成 bbox 解码再输出 4 个坐标有的版本只输出4 * reg_max的原始分布参数后处理再解码。无论哪种情况拆分逻辑都在最终输出张量上进行核心思路不受影响。2.2 双头输出的两种常见形态双头方案输出张量 1输出张量 2典型用途回归/分类分离(1, 8400, 4) 边界框(1, 8400, nc) 分类分数简化后处理量化调优更灵活检测/特征分离(1, 8400, 4nc) 检测结果(1, 8400, dim) ReID 或关键点特征目标跟踪、多任务学习本文主要演示第一种形态也就是把边界框和分类分数拆开。第二种形态是在第一种基础上的扩展如果后面要做 ByteTrack、StrongSORT 这类跟踪算法通常会在检测头旁边再加一个 ReID 特征头双头结构就是多任务扩展的基础。2.3 双头输出在 RK3588 部署中的实际价值从工程角度看双头输出在 RKNN 部署中有三个直接收益。第一量化更可控。默认单头输出的通道数是4 nc在 COCO 模型里通常是 84 个通道这些通道的取值范围差异很大回归通道的数值分布和分类通道完全不同。把输出拆成 4 通道和 nc 通道两个张量后每个通道组内部数值分布更集中RKNN 做 int8 量化时更容易找到合适的量化尺度。这是部署调优时的常见经验具体提升幅度要以量化后的精度对比为准。第二后处理更直观。板端拿到两个独立张量后直接对分类张量做argmax对回归张量做坐标解析不需要再依赖 torch 或 onnx 里的 slice 节点。C 后端如果直接解析 RKNN 输出这种张量边界清晰的格式写起来要舒服很多。第三便于扩展。双头输出本质上就是多任务输出的雏形。模型后面如果要同时输出检测框和行人 ReID 特征或者同时输出检测框和分割掩码都是在现有输出的基础上继续加分支。先把单头改成双头后面的扩展路径就顺畅了。3. 适用场景与使用边界这个改造方案适合以下四类开发者。第一类是正在 RK3588 上部署 YOLOv8发现量化后精度波动比较大的人。把输出层拆成双头再做量化是值得尝试的调优手段之一。第二类是准备把 YOLOv8 接到自研推理框架的工程师无论是 Python 后端还是 C 后端独立的输出张量更容易对接。第三类是做多任务的开发者比如检测加跟踪、检测加关键点、检测加分割双头输出是这些方向的基础结构。第四类是教学和模型结构研究场景想彻底理解 YOLOv8 检测头数据流的人。也有一些场景不适合这个改动。如果只是用官方权重在 x86 机器上跑 demo不需要改结构。如果模型已经在 RKNN 上稳定运行精度也满足要求没有必要为了结构上的“整洁”去冒险。如果团队没有能力做量化精度回归测试随意改输出层反而可能引入新问题。合规边界这里要单独提醒用于训练的公开数据集要确认授权范围不要直接拿未授权数据集微调或发布如果检测对象是人脸、行人、车牌等隐私敏感内容部署前要确认使用场景符合当地法律法规商用项目在正式发布前要重新用真实环境数据验证效果并保留测试记录。双头输出只是结构改造不改变模型本身的使用边界和责任边界。4. 环境准备与前置条件4.1 硬件准备RK3588 开发板一台是必须的建议搭配主动散热风扇和稳定的电源供电。RK3588 在做 NPU 推理和视频编解码时发热明显散热条件差会导致频率下降甚至系统不稳定。开发 PC 一台用于训练、导出 ONNX 和运行 RKNN-Toolkit2。RKNN 转换过程在 CPU 上可以完成但会占用较多内存建议开发机内存不低于 16GB。如果涉及重新训练模型还是建议使用 NVIDIA GPU训练效率和 CPU 完全不是一个量级。4.2 软件环境软件组件用途备注Ubuntu 18.04 / 20.04 / 22.04开发环境Windows 也可以用但 Ubuntu 更稳妥Python 3.8 及以上运行脚本具体版本以 rknn-toolkit2 要求为准PyTorch加载 YOLOv8 模型并导出 ONNX与 ultralytics 版本匹配Ultralytics YOLOv8 源码修改 Detect 头pip 安装或 GitHub 拉取源码均可ONNX检查输出节点配合 onnxruntime 做导出验证rknn-toolkit2PC 端模型转换需要安装到开发 PCrknn-toolkit-lite 或 RKNN C API板端推理板端运行时不装完整转换工具4.3 是否需要 GPU整个流程里模型训练建议 GPUONNX 导出在 CPU 上也能跑RKNN 转换和量化在 CPU 上可以执行但耗时会偏长板端推理只依赖 RK3588 内置 NPU不关联 GPU。所以如果你的开发机没有独立显卡用 YOLOv8s 这类小模型也能把这套流程完整跑通只是训练环节会比较吃力。对于只想验证双头改造流程的读者完全可以用官方预训练权重跳过训练环节。5. 实战前置拆解 YOLOv8 输出层源码进入改造之前先确认源码位置。以 Ultralytics 8.x 的常见目录结构为例检测头定义在ultralytics/nn/modules/head.py中核心类是Detect。这个类负责三个尺度检测头的输出拼接。Detect.forward里最关键的一行是把回归分支和分类分支拼接起来x[i] torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)cv2输出的是边界框回归相关特征通道数是4 * reg_max默认reg_max 16。cv3输出的是分类特征通道数是nc。拼接后每个尺度的特征图通道数变成4 * reg_max nc。之后三个尺度的特征图会被展平并合并。以 640x640 输入为例P3、P4、P5 特征图大小分别是 80x80、40x40、20x20展平后 anchor 数量为 6400、1600、400总计 8400。因此最终输出的第二个维度是 8400第三个维度是4 * reg_max nc。在实际输出前如果当前模型版本在 forward 里做了 DFL 解码那么最终输出的第三个维度可能是4 nc。判断标准很简单用torch.onnx.export导出后直接看输出张量的最后一维是 84COCO nc80就是已经解码过的坐标加分类是 6464 4*16纯 DFL 分布参数还要再解码一次。写后处理代码之前这个问题必须确认。6. 单头输出改成双头输出三种改法6.1 改法一直接修改 Detect forward这个改法的核心是保留推理结构不变只把最终 return 前的拼接逻辑替换成拆分逻辑。# 在 Detect.forward 中把默认的 return 之前加入拆分 box_feat, cls_feat torch.split(x_cat, (self.reg_max * 4, self.nc), dim1) box_out box_feat.transpose(1, 2) cls_out cls_feat.transpose(1, 2) if not self.dual_output: return torch.cat((box_out, cls_out), 2) return box_out, cls_out在类初始化时增加一个开关变量self.dual_output True这个开关的好处是可以在不改变训练逻辑的情况下通过切换开关分别导出单头和双头模型。训练时仍然保持trainingTrue走原来的分支推理时根据开关决定返回单头还是双头。需要注意不同版本的 Ultralytics 源码在forward里是否已经完成 DFL 解码存在差异。如果你的版本在拼接前已经调用过self.dfl则实际坐标维度是 4 而不是4 * reg_maxtorch.split的切分位置要相应改成 4。建议在修改前先打印x_cat.shape以实际输出为准。6.2 改法二导出 ONNX 时拆分输出节点如果不想改动 Ultralytics 源码可以在导出环节包一层 wrapper拦截模型输出并做拆分。这种方式适合快速验证不需要维护一份修改过的源码副本。import torch import torch.nn as nn class DetectOutputWrapper(nn.Module): def __init__(self, model): super().__init__() self.model model def forward(self, x): out self.model(x) # out 默认是 (1, 8400, 4 nc) 或 (1, 8400, 4 * reg_max nc) # 按实际最后一维切分 box, cls torch.split(out, (4, out.shape[-1] - 4), dim2) return box, cls wrapped_model DetectOutputWrapper(model) wrapped_model.eval() dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( wrapped_model, dummy_input, yolov8s_double.onnx, opset_version12, input_names[images], output_names[box_head, cls_head], dynamic_axes{images: {0: batch}, box_head: {0: batch}, cls_head: {0: batch}}, )wrapper 做法的改动面最小但它有一个明显限制如果原始模型输出在导出时经过了 DFL 解码和坐标变换wrapper 拿到的是模型最终输出拆分后的box_out已经是解码坐标。后续 RKNN 转换和后处理都会基于这个已解码的张量灵活性比改源码的方式低一些。6.3 改法三扩展为真正的多任务双头第三种改法适合检测加跟踪、检测加关键点这类多任务场景。此时两个输出头不是简单的拆分而是两个独立的分支。class MultiTaskHead(nn.Module): def __init__(self, det_head, feat_dim256): super().__init__() self.det_head det_head self.feat_branch nn.Sequential( nn.Conv2d(256, 128, 3, padding1), nn.ReLU(inplaceTrue), nn.Conv2d(128, feat_dim, 1), ) def forward(self, features): # features 是 Neck 输出的多尺度特征 det_out self.det_head(features) # 取 P5 特征作为 ReID 特征分支输入 feat self.feat_branch(features[-1]) return det_out, feat多任务头的改动量会明显增加因为它不仅涉及网络结构还涉及损失函数、正样本分配策略和特征对齐方式。对于只需要把单头拆成双头的场景不建议直接跳到这里。6.4 三种改法如何选择改法改动范围适合场景风险修改 Detect forward源码层影响训练验证导出全流程一套代码统一使用双头后续升级 Ultralytics 需要维护改动导出时拆分只改导出脚本不动源码快速验证结构可行性导出后的模型结构需重新检查多任务扩展新增模块改动训练和后处理检测加跟踪、加关键点等后处理复杂度明显上升我的建议是从改法二开始先花半小时把双头 ONNX 导出来确认输出节点和 shape 符合预期。确认可行之后再决定是否回到源码采用改法一。这样即使后续出现问题排查范围也更小。7. 导出 ONNX 并确认输出节点无论使用哪种改法最终都要落到 ONNX 模型上。导出这一步的目标很明确得到一个包含两个输出节点的 ONNX 文件节点名称稳定输出 shape 明确。# 导出示例实际参数以你的模型和源码为准 python export_yolov8_double.py --weights yolov8s.pt --imgsz 640 --opset 12如果采用改法二导出脚本里已经通过output_names指定了box_head和cls_head。如果采用改法一也要在自定义导出脚本里加同样的参数避免 ONNX 默认生成output_0、output_1这样的名字。导出完成后用 onnx 库检查输出节点import onnx model onnx.load(yolov8s_double.onnx) for out in model.graph.output: print(out.name) shape [d.dim_value for d in out.type.tensor_type.shape.dim] print(shape)预期看到两个输出节点名称是box_head和cls_headshape 第一维是动态 batch 或固定 1第二维是 8400第三维分别是 4 和 nc。如果只打印出一个输出节点说明 export 阶段又走了单头拼接分支需要回到 6.1 检查dual_output开关是否开启或者 wrapper 是否真的返回了两个张量。还有一种情况是 PyTorch 自动把两个输出合并了这通常发生在模型的forward返回值没有被正确解包时。8. RKNN 模型转换与量化配置ONNX 导出的问题解决之后进入 RKNN 转换阶段。这里以 rknn-toolkit2 的通用 API 为例具体参数名和函数签名会随工具链版本变化使用前先确认当前版本对应的接口。8.1 转换主流程from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[0, 0, 0]], std_values[[255, 255, 255]], target_platformrk3588, ) rknn.load_onnx(modelyolov8s_double.onnx) rknn.build(do_quantizationTrue, dataset./dataset.txt) rknn.export_rknn(yolov8s_double.rknn)mean_values和std_values必须与训练时的预处理一致不要照抄这里的 0 和 255否则量化后精度会受明显影响。常见做法是mean0, std255也就是输入像素归一化到 0 到 1但你的模型训练时未必是这个预处理方案。8.2 指定输出节点在部分 rknn-toolkit2 版本中load_onnx支持通过outputs参数指定希望保留的输出节点rknn.load_onnx( modelyolov8s_double.onnx, outputs[box_head, cls_head], )如果你使用的版本不支持这个参数更稳妥的办法是在 ONNX 层面就把输出节点处理好让 RKNN 默认保留 ONNX 的全部输出节点。转换完成后可以用rknn.list_outputs()查看当前模型实际保留的输出节点和输出顺序。8.3 量化数据集准备dataset.txt是量化阶段的数据列表每行写一张图片路径./images/img001.jpg ./images/img002.jpg ./images/img003.jpg量化数据集的选择对结果影响很大。图片数量不宜太少建议至少准备 5 到 20 张有代表性的图片覆盖目标尺度分布、光照变化和主要类别。如果检测场景是工厂缺陷检测量化数据集就要用工厂真实环境图片不能用 COCO 的通用图片代替否则部署后精度下降会很明显。建议先用do_quantizationFalse跑一遍非量化版本确认双头结构在 RKNN 里能正常转换和推理再开量化版本。这样能把结构问题从量化问题里分离出来。9. 板端推理与双头输出解析模型转换完成后把.rknn文件拷贝到 RK3588 开发板就可以开始板端验证。9.1 RKNN Python 推理示例import cv2 import numpy as np from rknnlite.api import RKNNLite rknn RKNNLite() rknn.load_rknn(yolov8s_double.rknn) rknn.init_runtime() img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img np.expand_dims(img, 0).astype(np.float32) outputs rknn.inference(inputs[img]) print(output count:, len(outputs)) print(box_head shape:, outputs[0].shape) print(cls_head shape:, outputs[1].shape)这里要注意输入预处理。如果rknn.config里已经设置了mean_values和std_valuesRKNN 在推理时会自动完成归一化此时不要在 Python 脚本里再除以 255否则等于做了两次归一化。不同 RKNN 版本对输入 dtype 和布局要求不同建议先跑通一个输入数据的最小示例再套用到实际检测图。运行后预期看到output count为 2box_head的 shape 是 (1, 8400, 4)cls_head的 shape 是 (1, 8400, nc)。如果output count是 1说明板端加载的模型仍是单头结构需要回到转换步骤检查是否导入了双头 ONNX。9.2 双头后处理逻辑双头输出的后处理流程比单头直观很多以 Python 示例说明def postprocess_double(box_out, cls_out, conf_thres0.25): # box_out: (1, 8400, 4) # cls_out: (1, 8400, nc) scores cls_out.max(axis-1) # (1, 8400) labels cls_out.argmax(axis-1) # (1, 8400) mask scores[0] conf_thres boxes box_out[0][mask] scores scores[0][mask] labels labels[0][mask] # 继续按类别做 NMS这里省略具体实现 keep nms_per_class(boxes, scores, labels, iou_thres0.45) return boxes[keep], scores[keep], labels[keep]单头后处理在解析时先要从(1, 8400, 4nc)里切出前 4 维和后 nc 维再做argmax和 NMS。双头后处理则少了一步 slice直接把两个张量交给不同的处理函数。对于 C 端来说这个差异更明显输出张量的内存布局是连续的取数逻辑更简洁。9.3 判断部署成功的标准部署是否成功可以从四个维度判断。第一输出节点数量正确得到两个独立张量。第二张量 shape 与 ONNX 导出阶段一致不能因为 RKNN 转换而多出或丢失维度。第三用同一张测试图对比板端推理结果与 PC 端 PyTorch 推理结果框的位置、类别、置信度排序趋势一致。第四连续推理多帧没有内存错误和异常退出耗时稳定在一个合理区间。需要说明的是板端推理耗时和显存占用受模型版本、输入分辨率、是否量化、NPU 频率设置共同影响这里不写死参考数值以你实际开发板的测试结果为准。10. 常见问题与排查方法问题现象可能原因排查方式解决方案导出 ONNX 后只有一个输出节点模型在导出时仍走了单头拼接分支打印 ONNX 输出节点名称检查dual_output开关或 wrapper 返回值输出节点名称不是box_head、cls_head导出时未指定output_names用 onnx 打印输出节点在torch.onnx.export里显式指定双头 shape 与预期不符torch.split切分维度算错打印x_cat.shape核对reg_max和nc的具体取值RKNN build 报算子不支持导出的 ONNX 带了多余后处理节点查看 RKNN 转换日志定位算子精简导出图只保留模型主体量化后精度明显下降量化数据集与真实场景差异大对比 float 与 int8 输出增加代表性图片调整 mean/std板端返回的输出顺序不确定输出顺序由 ONNX 节点顺序决定打印每个输出的 shape不要硬编码顺序通过 shape 判断板端推理耗时异常高NPU 没有正确启用或推理走了 CPU查看 NPU 使用状态和日志确认 runtime 初始化参数正确多路摄像头并发时程序卡顿多个线程共用一个 RKNN runtime加日志观察线程阻塞位置设计线程安全队列避免并发调用同一个 context排查时最有效的方法是逐步缩小范围。先确认 ONNX 结构正确再确认 RKNN 转换日志无异常最后确认板端输入预处理一致。绝大多数问题都出在这三个环节之间的数据假设不一致。11. 最佳实践与使用建议第一源码改法务必保留开关。在Detect类里加一个dual_output变量不要直接删掉单头分支。这样训练和导出可以灵活切回原结构排查问题也多一个对照选项。第二改完结构先在 PC 上用固定输入验证。构造一个随机输入分别跑原始模型和双头模型对比单头输出在切分后与双头输出的数值是否一致。这一步能拦截掉 80% 的结构错误。dummy torch.randn(1, 3, 640, 640) with torch.no_grad(): out_single model_single(dummy) out_dual model_dual(dummy) # 对比单头输出的前 4 列与双头 box_out 是否一致 assert torch.allclose(out_single[0, :, :4], out_dual[0][0, :, :4], atol1e-5) # 对比单头输出的后 nc 列与双头 cls_out 是否一致 assert torch.allclose(out_single[0, :, 4:], out_dual[1][0, :, :], atol1e-5)第三导出 ONNX 后立刻用 onnxruntime 检查输出节点和 shape。这一步成本最低发现节点数量不对直接回改源码不要等到 RKNN 转换阶段才排查。第四转换 RKNN 时先跑非量化版本。do_quantizationFalse跑通结构后再打开量化并准备数据集。把结构和精度分成两个阶段验证排查效率会高很多。第五量化数据集要和真实业务场景分布一致。数据集图片要覆盖远近目标、不同光照、不同类别数量和代表性比图片分辨率更重要。第六板端接口服务要做好访问控制。通过 RKNN Python API 或 C API 封装成服务后只监听内网地址加上简单的鉴权避免未授权调用消耗 NPU 资源。第七合规使用。模型权重、训练数据集、测试图片都要确认来源和授权。涉及人脸、车辆、行人等敏感场景部署前要评估隐私风险商用前必须做效果复核并留存测试记录。第八多摄像头场景要设计任务队列。多个视频流共用同一个 NPU 时不要让每个线程都创建独立的 RKNN runtime。推荐的做法是维护一个全局推理队列输入帧统一入队推理线程串行消费结果再分发回对应摄像头线程。这样可以充分利用 NPU同时避免并发访问导致的上下文切换开销。12. 总结与后续方向这次改造的核心可以用一句话概括在 YOLOv8 的 Detect 输出层里把原本拼接的边界框分支和分类分支拆成两个独立张量然后让 RKNN 保留这两个输出节点。整个链路是源码修改到 ONNX 导出再到 RKNN 转换最后到板端后处理。每步都有一个明确的验证动作源码改完对比张量数值ONNX 导出后检查节点数量和 shapeRKNN 转换后先跑非量化版本板端推理后对比单头模型的后处理结果。只要这四步验证都通过双头输出改造就算彻底跑通。下一步可以做的扩展方向很多。如果你主要做小目标检测可以仿照这个思路在输出层加一个高分辨率小目标头如果要做目标跟踪可以在双头基础上加 ReID 特征分支如果要做缺陷检测和视频监控还可以把 RK3588 的多路视频解码能力和这个双头模型接到同一个流水线里形成多摄像头并发检测系统。把这套输出层修改流程跑通后后面加任何自定义输出头思路都是一样的。