YOLOv8部署RK3588:单头输出改双头输出的完整实践

📅 发布时间:2026/9/8 8:35:11
YOLOv8部署RK3588:单头输出改双头输出的完整实践
把 YOLOv8 部署到 RK3588 上最难的一步往往不是 rknn 转换本身而是拆解输出层。很多人按教程导出 ONNX、用 RKNN-Toolkit2 转成 rknn 模型后发现推理接口只返回一个像[1, 84, 8400]的大张量也就是单头输出。这个张量里既有边界框回归结果又有 80 类分类结果想在板端单独拿坐标或单独拿类别分数必须按通道偏移去切代码可读性和维护性都很差。这篇文章会围绕“单头输出怎么变成双头输出”展开先拆解 YOLOv8 输出层的内部结构再给出两种可落地的改造方案最后在 RK3588 上完成 RKNN 转换、板端推理和后处理验证。适合已经在 PC 上训练过 YOLOv8、现在想把模型落地到 RK3588 或 RKNN 平台的开发者。1. YOLOv8 输出层拆解单头输出到底是什么1.1 从检测头说起分类分支和回归分支为什么会被合并YOLOv8 是一个 anchor-free 检测器它的检测头采用了解耦设计也就是 Detect 模块内部同时存在两条分支。回归分支 reg 负责预测每个 anchor 位置的边界框坐标分类分支 cls 负责预测该位置属于每个类别的概率。这两条分支在训练阶段是分开计算损失的这样可以让模型分别关注“框得准不准”和“分得对不对”。但到了推理导出阶段很多版本的 Ultralytics 源码会把两条分支的结果在通道维度上拼接成一个张量形成单头输出。这样做的原因很直接导出 ONNX 时输出节点越少部署工具链越容易处理后端 NMS 也可以只遍历一个张量拿到前 4 个通道作为坐标、后面 80 个通道作为分类分数。这里“单头输出”指的不是 CPU 核数或者检测头的数量而是 ONNX 模型在输出层只有一个输出分支。理解这一点后面拆双头时就不会晕。1.2 一个 640 输入模型输出张量每一维代表什么以一个常见的 YOLOv8n 模型为例输入图片尺寸是 640×640类别数按 COCO 的 80 类计算。经过 Backbone 和 Neck 之后模型会在三个不同尺度上输出检测结果P3 特征图尺寸 80×80对应 stride 8展平后有 6400 个位置。P4 特征图尺寸 40×40对应 stride 16展平后有 1600 个位置。P5 特征图尺寸 20×20对应 stride 32展平后有 400 个位置。三个尺度展平后的位置总数是 6400 1600 400 8400。因此单头输出张量的形状是[1, 84, 8400]。这里的 84 等于 4 加 80前 4 个通道是边界框坐标后 80 个通道是分类分数。有的部署代码会习惯把张量转成[1, 8400, 84]第二种排列方式只是把每个位置的所有通道放在连续内存里更接近 NMS 遍历时的习惯。输出形态输出数量输出 shape 示例说明单头输出1[1, 84, 8400]坐标和分类拼接在一个张量里双头输出2[1, 4, 8400]和[1, 80, 8400]回归分支和分类分支独立输出三尺度输出3[1, 84, 6400]、[1, 84, 1600]、[1, 84, 400]每个尺度单独输出便于按 stride 处理1.3 单头、双头、三尺度输出三种形态怎么选单头输出适合最常规的部署方式把后处理逻辑固定成“先取前 4 通道再取后面 N 个通道”代码简单但灵活性差。双头输出是把分类和回归彻底分开适合想要自定义 NMS、按类别做阈值过滤、或者单独调试定位分支和分类分支的场景。拆开之后分类分支可以先做 sigmoid 和阈值过滤回归分支再做坐标还原两条链路互不干扰。三尺度输出则是按特征图尺度拆开的方案。如果项目里对小目标检测做了额外处理或者希望针对不同尺度使用不同置信度阈值建议直接按尺度输出。需要说明的是双头输出并不比单头输出更高明它只是把原本在板端后处理里要做的“切分”动作提前到模型导出阶段完成。具体怎么选取决于后处理代码的维护方式。2. RK3588 部署为什么不能一直用单头输出2.1 RKNN 后处理与输出节点有直接关系RKNN 推理时rknn.inference()返回的输出顺序和 ONNX 模型的输出节点顺序一致。如果 ONNX 只有一个输出节点那么板端拿到的就是一个大张量。要在 C 代码里解析单头输出需要按固定步长从一维缓冲区里取数下标稍微算错坐标、置信度、类别就会整体错位。拆成双头之后RKNN 会返回两个独立的输出 buffer。分类分支可以单独传给 topk 模块回归分支单独做坐标缩放后处理代码的可读性会好很多。2.2 单头输出在板端会遇到的具体问题在板端后处理中单头输出最容易出现的问题是通道偏移。[1, 84, 8400]的布局决定了每个 anchor 对应的数据是 84 个连续浮点数前 4 个是坐标后 80 个是分类分。如果代码想单独遍历所有 anchor 的分类分数就需要不断跳过坐标维缓存命中率低循环也不够直观。更麻烦的是一旦项目改了类别数例如从 80 类改成 2 类单头输出的通道数会从 84 变成 6。所有依赖“84 这个数字”的硬编码后处理代码全部要改一遍。双头输出可以降低这种耦合reg 分支始终是 4 通道cls 分支只随类别数变化。2.3 双头输出能省掉哪些中间操作单头输出在后处理时需要先用memcpy或循环把坐标部分和分类部分分别拷贝出来再做后续处理。双头输出相当于把这个动作从“运行期”提前到“模型导出期”。实际省掉的代码量不大省掉的是对数据布局的假设。如果你的板端后处理已经写好了类似这样的伪代码for (int i 0; i 8400; i) { float *data output i * 84; float x1 data[0]; float y1 data[1]; float x2 data[2]; float y2 data[3]; float score data[4 class_id]; }那整个循环强依赖 84 通道连续排列。改成双头后reg 数组和 cls 数组可以分别读取后处理逻辑也更清晰。2.4 什么时候应该保留单头如果只是跑 rknn_model_zoo 里的标准示例或者后处理代码已经稳定运行并且没有调整检测头结构的计划那就不要为了拆而拆。增加输出节点意味着 RKNN 转换时输出分支变多板端也需要多管理一个 buffer量化后需要验证的数值关系也更多。对比维度单头输出双头输出输出节点数12后处理切分在板端按通道偏移切导模型时已切好代码耦合度强依赖 4 num_classes低类别数变更需要改板端代码cls 分支 shape 变化逻辑不变扩展小目标头不利有利转换和板端内存开销较低略高3. 环境准备PC 端导出、RKNN 转换和板端推理的分工3.1 硬件和软件环境清单RK3588 部署链路通常由两部分组成PC 负责模型导出和 RKNN 转换开发板负责实际推理。PC 端建议使用 x86_64 Linux 环境开发板可以是常见的 RK3588 开发板例如正点原子等品牌。环境用途建议PC导出 ONNX、修改输出层、RKNN 转换Linux 系统安装 Python 3.8 左右版本RK3588 开发板板端推理、后处理验证系统使用 Ubuntu 或 Debian确认 RKNPU 驱动正常模型文件待部署的 YOLOv8 模型自己训练的 .pt 或官方预训练模型均可推理输入测试图片或视频流先准备若干张与训练分布接近的图片需要注意一个版本问题RKNN-Toolkit2 和板端 RKNPU 驱动的版本需要匹配。安装 RKNN-Toolkit2 前先确认开发板烧录的系统镜像里 RKNPU 驱动版本再选择对应的工具包。代码中出现的关键参数在不同版本里可能有差异落地前要以上机环境为准。3.2 RKNN-Toolkit2 安装与初始化RKNN-Toolkit2 通常通过 wheel 包方式安装不建议直接从公开 PyPI 拉取因为工具链版本与芯片驱动绑定。安装完成后先跑一个最小初始化验证python -c from rknn.api import RKNN; print(rknn ok)如果导入失败先检查 Python 版本、依赖库版本和 wheel 包是否下载完整。3.3 准备 YOLOv8 模型并确认输出节点先准备一个可用的 YOLOv8 ONNX 模型。如果还没有训练好的模型可以使用官方预训练权重导出yolo export modelyolov8n.pt formatonnx imgsz640导出后用 Netron 或 onnx 脚本查看输出节点确认当前模型的输出形态。import onnx model onnx.load(yolov8n.onnx) for i, out in enumerate(model.graph.output): print(i, out.name)这里要特别提醒不同 Ultralytics 版本导出的输出名可能不同后面做图修改时不要硬编码输出名要先打印确认。4. 把单头输出改成双头输出两种可落地方案4.1 方案一修改 Detect head 直接导出双输出Ultralytics 的 Detect head 在推理导出时会执行类似下面的逻辑先对回归分支做 DFL 解码再把坐标和分类分支拼接。不同版本的源码结构不一样但核心思路是一致的。在ultralytics/nn/modules/head.py中找到 Detect 类的forward方法定位到导出分支。以某个常见的 8.x 版本为例导出时会出现类似这样的拼接逻辑y torch.cat((dbox, cls.sigmoid()), 1) return y if self.export else (y, x)要改成双头输出可以把这个导出分支修改为返回两个独立的张量if self.export: return dbox, cls修改时要注意两个关键点。第一dbox必须已经经过 DFL 解码否则 reg 分支输出的是 64 通道的分布参数而不是 4 通道的坐标。第二不要在导出前提前对cls做 sigmoid把这个动作留给板端后处理保持通用性。修改后重新导出yolo export modelyolov8n.pt formatonnx imgsz640再用 onnx 脚本检查输出节点。这个方案的优点是可以从源头得到干净的双头 ONNX缺点是升级 Ultralytics 版本时源码改动容易被覆盖维护成本高。4.2 方案二在不改源码的情况下用 ONNX 脚本拆输出如果不想动训练框架源码或者已经生成好了单头 ONNX可以直接用 onnx 库修改计算图。核心思路是单头输出是坐标和分类拼接后的结果。查找图里的 Concat 节点确认 concat 的两个输入分别是坐标分支和分类分支。找到后把 concat 的两个输入分别作为新输出即可。下面这段代码是思路示例实际运行前需要先确认自己模型里 Concat 节点的名称和输入顺序import onnx from onnx import helper model onnx.load(yolov8n.onnx) graph model.graph # 1. 找到输出节点前面的 Concat concat_node None for node in graph.node: if node.op_type Concat: # 根据输出 shape 判断是否是要找的目标或者直接打印节点名 print(node.name, node.output) concat_node node break if concat_node is None: raise RuntimeError(no concat node found) # 2. 确认 concat 的两个输入并给它们接上 Identity 输出 for idx, input_name in enumerate(concat_node.input): out_name fhead_{reg if idx 0 else cls}_output identity helper.make_node( Identity, inputs[input_name], outputs[out_name], namefsplit_head_{idx} ) graph.node.insert(0, identity) new_output helper.make_tensor_value_info( out_name, onnx.TensorProto.FLOAT, None ) graph.output.append(new_output) # 3. 保存新模型 onnx.checker.check_model(model) onnx.save(model, yolov8n_dual_onpaper.onnx)这段代码会把 Concat 的两个输入导出为两个输出节点同时保留原来的 Concat 输出。如果不想要原来的单头输出还要把旧的 graph.output 移除。不管用哪种方式修改后都要用 onnxruntime 或其他推理框架跑一遍确认输出 shape 是否符合预期。如果图里没有明显的 Concat 节点也可以用 Slice 方式从[1, 84, 8400]上切出[1, 4, 8400]和[1, 80, 8400]。这个思路更通用但要在图里插入 Slice 算子并添加对应的 starts、ends、axes 常量代码量会更多。除非 Concat 查找失败否则优先推荐 Identity 方案。4.3 两种方案对比和后处理适配对比维度修改 Detect headONNX 脚本修改是否改训练源码是否对 Ultralytics 版本敏感度高中输出稳定性高依赖图匹配逻辑适合场景长期维护自己的检测工程快速在部署流水线里处理已有模型风险点升级框架后源码冲突输出节点名和 Concat 输入顺序变化无论用哪种方式最后都要确认双头输出的语义reg 输出是坐标cls 输出是未经过 sigmoid 的分类分数。不要把坐标分支和分类分支搞反。注意修改 ONNX 输出层后不要急着跑 RKNN 转换先在 PC 上用 onnxruntime 跑一张测试图打印两个输出的 shape确认 shape 是[1, 4, 8400]和[1, 80, 8400]再进入下一步。5. RKNN 转换、板端推理与结果验证5.1 用 RKNN-Toolkit2 转换双头 ONNX得到双头 ONNX 后开始做 RKNN 转换。以下代码在 PC 上运行from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[0, 0, 0]], std_values[[255, 255, 255]], target_platformrk3588 ) if rknn.load_onnx(modelyolov8n_dual.onnx) ! 0: raise RuntimeError(load onnx failed) if rknn.build(do_quantizationTrue, datasetdataset.txt) ! 0: raise RuntimeError(build failed) rknn.export_rknn(yolov8n_dual.rknn)dataset.txt 里每行写一张校准图片的路径不需要做标注只用于量化时的激活值统计。图片要经过和真实推理一致的前处理否则量化精度会受影响。如果只是一个功能验证可以先关闭量化rknn.build(do_quantizationFalse, datasetdataset.txt)关闭量化能得到更接近浮点精度的结果方便排查输出层改造是否出错。5.2 板端推理时如何拿到两个输出在 RK3588 板端运行时代码结构和 PC 上类似只是 init_runtime 时传入板端信息。先做一个最简单的推理验证import cv2 import numpy as np from rknn.api import RKNN rknn RKNN() rknn.load_rknn(yolov8n_dual.rknn) rknn.init_runtime(targetNone) img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img img.astype(np.float32) outputs rknn.inference(inputs[img]) print(output num:, len(outputs)) print(reg shape:, outputs[0].shape) print(cls shape:, outputs[1].shape)正常情况下输出是两个张量shape 分别接近[1, 4, 8400]和[1, 80, 8400]。如果只输出一个张量说明前面的 ONNX 修改没有成功需要回退到第 4 章检查 ONNX 输出节点。5.3 双头输出的后处理最小流程拿到双头输出后后处理逻辑可以拆成两条链路。先用 Python 写最小实现验证reg, cls outputs # [1, 4, 8400], [1, 80, 8400] cls cls[0].T # [8400, 80] scores cls.max(axis1) class_ids cls.argmax(axis1) conf_mask scores 0.25 selected_scores scores[conf_mask] selected_ids class_ids[conf_mask] selected_boxes reg[0].T[conf_mask] # [n, 4]这里有一个容易搞混的地方。YOLOv8 的坐标解码方式在不同版本里可能是 xyxy、ltrb 或 cxcywh。如果 ONNX 导出时走了 Detect head 里的 decode_bboxes那么 reg 分支输出的通常是相对于特征图的坐标需要乘以对应 stride 才能映射回原图。实际项目里必须根据你导出的那个模型打印一组输出和单头模型的结果对比后再确定坐标还原公式。后续 NMS 部分可以选择普通遍历或按类别分组遍历双头输出的好处是可以先做一次 topk 过滤只保留少量高分候选框减少 NMS 的计算量。5.4 验证结果应该长什么样双头输出改造是否成功可以从两个层面验证。第一层是数值验证。用同一张输入图分别推理单头模型和双头模型然后把双头模型的 reg 和 cls 在通道维拼接回来理论上应该近似等于单头模型的输出。量化后两者会有偏差但坐标和高置信度类别的趋势应当一致。第二层是效果验证。把双头输出的检测框画到原图上和单头模型的结果对比。正常情况下同一张图上的检测框数量和位置应该基本一致。如果框的位置完全错乱优先检查坐标是否没有乘 stride或者 reg 分支是否没有经过 DFL 解码。注意量化模型和浮点模型的输出不能要求完全一致。只要目标框的重叠率在合理范围内分类置信度的排序没有大变化就可以认为输出层改造是成功的。6. 拆输出层时最容易踩的 5 个坑6.1 分不清 DFL 分布和解码后的坐标现象修改 Detect head 后reg 输出 shape 不是[1, 4, 8400]而是[1, 64, 8400]。原因YOLOv8 的 DFL 模块把每个坐标编码成 16 个 bin 的分布4 个坐标加起来是 64 通道。如果导出时没有执行 DFL 解码reg 分支输出的就是分布参数不能直接当坐标用。检查方式打印 reg 输出的 shape看到 64 就说明没有 decode。处理方式在导出分支里先调用 decode_bboxes把 64 通道变成 4 通道再输出。6.2 混淆 anchors 排列顺序和通道顺序现象拆成双头后后处理能跑通但画框位置错位。原因8400 个位置不是按乱序排的而是 P3、P4、P5 三个尺度展平后按顺序拼接。P3 是前 6400 个位置P4 是第 6400 到 8000 个位置P5 是最后 400 个位置。很多代码直接从第 0 个位置开始按 stride 8 解析结果把后面尺度的坐标映射错了。处理方式先在 ONNX 模型上打印 3 个尺度对应 anchor 的区间再写后处理。小目标检测头加入后anchor 总数和尺度顺序都会变化不能沿用老代码。6.3 漏掉 sigmoid 或做了两次 sigmoid现象检测置信度整体偏高或偏低框不可信。原因分类分支的原始输出是 logits需要做 sigmoid 转成概率。如果把已经做过 sigmoid 的 cls 输出当成 logits后处理又做一次 sigmoid会让分数整体压低反过来漏掉 sigmoid分数会超过 1 或者变成负值。处理方式约定好 sigmoid 只做一次。建议在板端后处理里做 sigmoid导出 ONNX 时保持 cls 是原始 logits这样模型输出更通用。6.4 在 RKNN 里加错 permute 导致维度错位现象打印 shape 正确但读取数值时配对错乱。原因RKNN-Toolkit2 对 ONNX 的输入输出布局有默认处理。有的项目为了让 NMS 好写在 ONNX 里加了 transpose把通道维挪到最后。如果 RKNN 配置里又做了一次类似 permute 的操作就会多转一次。处理方式固定 ONNX 里的布局转换前在 PC 上用 onnxruntime 确认输出布局再决定 RKNN 后处理里是否需要转置不要在两处各转一次。6.5 量化校准数据与输出尺寸不匹配现象int8 量化后双头输出的分类分数波动很大检测框不稳定。原因dataset.txt 里的校准图片直接读了原图没有做 640×640 的缩放和通道转换。校准图片和真实推理输入差异太大导致量化时激活值统计不准。处理方式校准图片要经过和推理完全一致的前处理流程包括 resize、通道顺序、归一化参数。校准集不需要很大20 到 100 张有代表性的图片通常就够。问题现象常见原因检查方式处理建议reg 输出 64 通道DFL 未解码打印输出 shape导出时执行 decode_bboxes画框位置错乱anchor 尺度顺序混淆对照 8400 区间按 P3、P4、P5 顺序写后处理置信度异常sigmoid 重复或漏做打印 cls 数值范围只保留一处 sigmoidshape 对但数值乱布局重复转发对比单头模型输出统一 ONNX 布局避免二次 permute量化后精度下降明显校准集前处理不一致检查 dataset.txt 图片使用统一预处理流程7. 从教程到量产RK3588 上输出层设计的最佳实践7.1 输出形态的选型建议在真正量产或长期维护的项目中输出形态的选择应该由后处理代码的维护方式决定而不是由“哪种听起来更高级”决定。如果板端后处理已经稳定且不希望增加输出 buffer 数量保留单头输出是合理的。如果项目需要经常调整类别数、增加小目标检测头、或对分类阈值有细粒度控制建议改成双头输出。如果希望能按不同尺度独立调参可以使用三尺度输出但每个尺度都要独立解码代码量会明显增加。项目场景推荐输出形态理由官方 demo 快速验证单头代码简单示例多多类别自适应项目双头类别变化只影响 cls 分支小目标检测优化双头或多尺度便于单独调试特征层纯 C 部署追求极致性能单头或双头按实际 buffer 开销测试后决定7.2 量化、校准和精度验证RK3588 的 NPU 对 int8 支持很好但量化精度取决于校准数据是否贴近真实场景。建议收集一批真实场景图片作为校准集而不是直接用训练集。转换后不要只检查单张图至少用 20 张以上测试图对比浮点模型和量化模型的目标框。如果量化后某一类目标出现明显漏检可以检查该类目标在图片中的像素占比必要时调整校准集构成。输出层拆成双头后reg 分支和 cls 分支会共用同一个量化参数数值敏感度不同需要一起观察。7.3 板端 C 后处理与 NPU 任务拆分RK3588 的优势在于 NPU 负责卷积计算CPU 负责后处理。双头输出在后处理时建议先做 cls 分支的 topk 过滤只保留分数较高的候选框再进入 NMS。如果直接对 8400×80 的全量结果做 NMSCPU 耗时会被放得很大。在多路视频流场景下推荐把采集、推理、后处理拆分到不同线程。NPU 推理结果放到共享缓冲池后处理线程只读结果避免推理和后处理互相阻塞。输出层越简单越容易把 CPU 资源留给其他系统控制任务。7.4 可复用的部署前检查清单以下清单可以直接复制到项目文档里每次部署新模型时逐项检查ONNX 输出节点数量是否符合预期打印 graph.output 确认。reg 分支 shape 是 4 通道不是 64 通道 DFL 分布。cls 分支输出的是 logits板端只做一次 sigmoid。anchor 拼接顺序已确认P3、P4、P5 区间清楚。用同一张图对比单头模型和双头模型的输出数值趋势一致。RKNN 转换前先做一次浮点 build验证输出节点和 shape。量化校准图片使用统一前处理不直接用训练集原图。板端推理打印len(outputs)和每个输出 shape。画框结果和单头模型对照位置和类别无明显错乱。检查 RKNN-Toolkit2 版本和板端 RKNPU 驱动版本匹配。拆解 YOLOv8 输出层是部署排错的基本功。单头变双头这件事表面上是改几行导出代码核心是能说清每一个维度代表什么、每个分支在什么位置被解码、后处理里应该在哪里做 sigmoid。建议先在 PC 上用 onnxruntime 把双头输出和原单头输出对齐确认数值没有异常再进入 RKNN 转换和板端开发。下一步可以继续研究 RKNN 的 zero_copy 接口、多路视频流推理以及把 NPU 推理和 CPU NMS 拆成流水线这些都是在 RK3588 上把 YOLOv8 从“能跑”推向“能稳定跑”的关键方向。