尧图网页设计 
在做 RK3588 端侧部署时很多朋友会遇到一个很尴尬的情况模型在电脑上跑得好好的精度也够一旦转到 RKNN 就出现各种算子不支持、后处理写法别扭、多任务输出没法拆分的问题。尤其是 YOLOv8 这种自带解耦头的模型默认输出结构在边缘侧并不总是“最优解”。本篇文章就围绕 RK3588 YOLOv8 这个组合从源码层面拆解 YOLOv8 输出层的设计再一步步把默认的单头输出改造成双头输出并完整走一遍训练、导出、RKNN 转换和板端部署流程。这篇文章适合已经跑通 RK3588 基础环境、正在做模型优化或想自定义 YOLOv8 检测头的开发者。零基础的朋友可以先补一下 RKNN Toolkit2 的基础用法再回来看本文的改动思路。1. 为什么要拆解输出层从单头到双头的实际场景1.1 输出层在 YOLOv8 中扮演的角色YOLOv8 的网络结构可以粗分成三块Backbone主干网络、Neck特征融合、Head检测头。Backbone 负责提取图像特征Neck 把不同尺度的特征做融合最后交给 Head 输出目标类别和边界框位置。输出层也就是 Head 部分直接决定了“网络最终吐出什么”。默认情况下YOLOv8 的 Head 会在三个尺度上分别输出预测结果。这三个尺度对应特征图的下采样倍数通常是 8 倍、16 倍、32 倍。小幅度的特征图负责检测大目标大幅度负责检测小目标。在 YOLOv8 的源码里这个 Head 是由Detect类实现的。它的核心逻辑不复杂对每个尺度的特征图做卷积预测出目标的类别概率和边界框坐标最后把多尺度的结果拼到一起形成最终输出。如果只是做标准的目标检测直接用默认 Head 没有任何问题。但一旦涉及到“检测 分类 关键点”这样的多任务场景或者想把小目标检测单独拆出来做优化默认的单一输出结构就不够灵活了。1.2 单头与双头的本质区别先解释一下本文语境下的“单头”和“双头”。默认 YOLOv8 的Detect类虽然在内部把分类和回归分支分开了但从外部看最终只输出一个融合后的结果这个叫单头输出。它的结构大致是这样的输入特征图 ├── 分类分支Conv Conv Conv2d └── 回归分支Conv Conv Conv2d ↓ 拼接成一个 Tensor 输出所谓双头输出就是把这个单一结果拆成两个独立的输出分支。比如输入特征图 ├── 输出头 A负责类别预测 └── 输出头 B负责边界框回归还可以是另一种双头在原本的检测头基础上额外增加一个“辅助头”比如一个专门输出小目标检测结果的浅层头或者一个输出关键点坐标的头。在 RK3588 上做双头改造常见动机有三个多任务输出一个头做检测一个头做属性分类或者一个头做检测一个头做关键点。小目标增强在原有 P3/P4/P5 三个输出尺度的基础上增加一个 P2 浅层输出头让模型对小目标更敏感。NPU 部署优化把输出拆开后RKNN 可以分别处理不同尺度的结果后处理代码更容易编写和维护。1.3 RK3588 部署场景下的真实需求RK3588 内置的 NPU 算力虽然不错但和桌面级 GPU 不同它对网络结构的兼容性有自己的脾气。很多在 PyTorch 里能跑的算子转到 RKNN 时可能就不支持或者需要特殊写法。我在实际项目中遇到的一个典型场景是需要同时检测目标并判断目标朝向。原来的做法是单独训练一个 YOLOv8 检测模型和一个分类模型部署时串行推理。这样做有两个问题一是耗时翻倍二是两套模型的预处理和后处理互相独立工程上很繁琐。改成双头输出后检测和朝向分类在一个模型里完成前处理只做一次NPU 推理也只做一次后处理从两个模型的结果里各自取数据即可。整体帧率提升很明显。另一个常见场景是小目标检测。标准 YOLOv8 的三个输出尺度对密集小目标效果有限。很多做工厂缺陷检测、无人机航拍目标检测的朋友都会尝试加一个 P2 输出头。这个改动其实也可以理解成一种“双头到多头”的扩展思路。所以学会从源码层面修改 YOLOv8 的输出层是做端侧模型定制绕不开的基本功。2. 读懂 YOLOv8 的检测头源码2.1 YOLOv8 的 Decoupled Head 结构YOLOv8 和 YOLOv5 一个很大的区别就是 Head 从耦合头改成了解耦头。YOLOv5 的检测头是直接在特征图上预测(4 1 nc) * 3个通道其中 4 是边界框坐标1 是置信度nc 是类别数3 是 anchor 数量。这种方式把分类和回归耦合在同一个输出通道上。YOLOv8 则把分类和回归拆成两个独立分支。在源码ultralytics/nn/modules/head.py的Detect类中可以看到两个核心卷积模块cv2回归分支输出4 * reg_max个通道cv3分类分支输出nc个通道这里的reg_max是 YOLOv8 度量学习中 DFLDistribution Focal Loss模块的参数默认值是 16。也就是说每个边界框的每条边不是直接回归一个坐标值而是回归 16 个离散概率再通过 DFL 层加权求和得到最终坐标。这种结构图大致如下输入特征图P3/P4/P5 ├── cv2 分支Conv - Conv - Conv2d(4 * reg_max) │ ↓ │ 4 * reg_max 通道的回归特征 └── cv3 分支Conv - Conv - Conv2d(nc) ↓ nc 通道的分类特征2.2 Detect 类的初始化与 Forward 流程打开 YOLOv8 源码Detect 类的初始化函数中比较关键的部分如下class Detect(nn.Module): def __init__(self, nc80, ch()): super().__init__() self.nc nc self.nl len(ch) self.reg_max 16 self.no nc self.reg_max * 4 self.stride torch.zeros(self.nl) c2, c3 max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100)) self.cv2 nn.ModuleList( nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch ) self.cv3 nn.ModuleList( nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch ) self.dfl DFL(self.reg_max)这里ch是 Neck 输出的三个特征图的通道数。cv2和cv3都是 ModuleList每个特征图对应一套自己的卷积参数。forward方法做的事可以拆成下面几步遍历三个尺度的特征图分别经过cv2和cv3。对分类结果做 sigmoid 激活。对回归结果通过 DFL 计算边界框坐标。按尺度拼接并统一转成推理格式输出。核心代码如下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 ... # 推理模式拼接所有尺度的输出并返回注意这个torch.cat它把回归分支和分类分支拼接在通道维度上。这就是“单头”输出的核心所在外部调用模型时拿到的是一个已经拼好的结果。2.3 输出 Tensor 的形状与含义推理模式下YOLOv8 输出的格式是[batch, num_anchors, 4 nc]。num_anchors是所有尺度特征图的锚点总数。以 640x640 输入为例P3 特征图是 80x80有 6400 个锚点P4 特征图是 40x40有 1600 个锚点P5 特征图是 20x20有 400 个锚点合计 8400 个锚点。每个锚点对应4 nc个数值前 4 个是边界框坐标归一化的 x, y, w, h后面 nc 个是类别概率。了解这个形状很重要因为改造成双头输出后前 4 个数值和后面 nc 个数值会被拆到两个不同的 Tensor 中后处理代码需要相应调整。3. 环境准备与代码基线3.1 硬件与软件环境本文的环境如下版本可以根据自己的项目情况调整项目版本开发板RK3588本文以正点原子 RK3588 为例NPU 工具链rknn-toolkit2训练环境Python 3.8 PyTorch 2.0YOLOv8 源码ultralytics 8.x模型导出格式ONNX - RKNN操作系统可以是 Ubuntu也可以直接在 RK3588 的板端系统上做推理验证。训练和导出建议放在 PC 上完成RK3588 只负责 RKNN 模型推理。3.2 准备 YOLOv8 源码与预训练权重首先克隆 ultralytics 源码。如果网络不方便也可以直接在 PyPI 安装pip install ultralytics不过由于我们需要改源码里的Detect类建议直接使用 GitHub 源码方式git clone https://github.com/ultralytics/ultralytics.git cd ultralytics pip install -e .下载一个预训练权重作为基线比如 YOLOv8swget https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s.pt这个权重是 COCO 80 类预训练的后面做双头改造时我们可以先用它跑通流程再用自己的数据集重新训练。3.3 快速验证当前输出写一个最简单的推理脚本确认当前单头输出的形状import torch from ultralytics import YOLO model YOLO(yolov8s.pt) model.eval() dummy_input torch.randn(1, 3, 640, 640) with torch.no_grad(): result model.predict(dummy_input) print(type(result))如果要看更底层的模型输出形状可以直接拿model.model的 forward 结果import torch from ultralytics.nn.tasks import DetectionModel model DetectionModel(yolov8s.yaml) model.eval() dummy_input torch.randn(1, 3, 640, 640) with torch.no_grad(): output model(dummy_input) print(len(output)) for i, x in enumerate(output): print(i, x.shape)这里输出的就是推理阶段模型拼接后的最终结果。一共一个 Tensor形状是[1, 8400, 84]其中 84 4 80。记住这个基线结果后面改成双头输出后比较差异就有依据了。4. 核心改造单头变双头输出4.1 设计双头结构在动手改代码之前先把目标说清楚。我们这次要实现的“双头”是把默认的检测头从“一个融合输出”改造成“两个独立输出”输出头 A边界框回归结果形状为[batch, 8400, 4]输出头 B分类结果形状为[batch, 8400, nc]后处理时先从两个输出里分别取所需数据再合并计算最终检测框。这样做的好处是后处理代码更清晰不需要在一大块连续内存里切分。任务解耦后可以单独对分类分支做蒸馏、量化或优化。如果后续要扩展第三个头比如关键点头结构上也很方便。4.2 修改 Detect 类yolov8 的模型定义中检测头是动态加载的。我们需要修改ultralytics/nn/modules/head.py里的Detect类。核心思路是在forward中不再把cv2和cv3的结果cat在一起而是分别返回。下面给出修改后的核心代码class Detect(nn.Module): def __init__(self, nc80, ch()): super().__init__() self.nc nc self.nl len(ch) self.reg_max 16 self.no nc self.reg_max * 4 self.stride torch.zeros(self.nl) c2, c3 max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100)) self.cv2 nn.ModuleList( nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch ) self.cv3 nn.ModuleList( nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch ) self.dfl DFL(self.reg_max) def forward(self, x): shape x[0].shape dfl_out [] cls_out [] for i in range(self.nl): reg_branch self.cv2[i](x[i]) cls_branch self.cv3[i](x[i]) # 回归分支解码 b, c, h, w reg_branch.shape reg_branch reg_branch.view(b, 4, self.reg_max, h * w).transpose(2, 3).reshape(b, 4, h * w) dfl_out.append(self.dfl(reg_branch)) # 分类分支 sigmoid cls_out.append(torch.sigmoid(cls_branch.reshape(b, self.nc, h * w))) if self.training: # 训练阶段返回回归原始特征和分类特征便于 loss 计算 return dfl_out, cls_out # 推理阶段分尺度转换为 xywh 格式 anchors, strides self.make_anchors(x, self.stride, 0.5) results [] for i in range(self.nl): d dfl_out[i].transpose(1, 2) c cls_out[i].transpose(1, 2) results.append((d, c)) return results这段代码里有几个关键点需要解释一下reg_branch要经过 DFL 解码才能得到边界框坐标。cls_branch需要做 sigmoid得到类别概率。训练阶段和推理阶段返回的格式不一样。训练阶段需要保留原始特征方便 loss 计算推理阶段输出解码后的坐标和类别概率方便后处理直接使用。当然实际工程中很多人不会直接在主源码里大改而是继承Detect写一个DetectDual类。这样改动更小也不容易影响原有逻辑。比如class DetectDual(Detect): def forward(self, x): # 重写 forward拆分输出 ...然后在模型的 yaml 配置里把 head 的类型改成DetectDual。如果你用的是自定义数据集还要注意nc改成自己的类别数。比如检测车辆和行人nc2。4.3 修改 Loss 计算逻辑输出结构变了loss 计算也必须跟着改。默认情况下YOLOv8 的 loss 变化在ultralytics/utils/loss.py文件里调用。原来的forward返回拼接结果loss 会从拼接结果里切分分类分支和回归分支。改成双头输出后loss 函数需要直接接收两个分支的结果。修改原则是分类 loss 只依赖分类分支输出回归 loss 只依赖回归分支输出。点开 v8DetectionLoss 类把原先从pred中切分pred_distri和pred_scores的部分改成直接接收双头结果。大致逻辑如下class v8DetectionLoss: def __init__(self, model): # 原有初始化逻辑保持不变 ... def __call__(self, preds, batch): # preds 变成了 (dfl_out, cls_out) pred_distri, pred_scores preds ...需要注意的是训练阶段我们返回的是未解码的回归特征loss 内部会通过 DFL 计算分布损失。保持和官方一致的流程即可。4.4 修改导出逻辑模型训练完以后要导出 ONNX。导出的逻辑在ultralytics/nn/tasks.py的Model.export或exporter.py里。默认导出时会把 forward 的输出拼接成一个维度输出。改双头后导出时会输出两个节点这符合我们的需求。导出命令yolo export modelruns/train/exp/weights/best.pt formatonnx opset12或者直接用 Python APIfrom ultralytics import YOLO model YOLO(runs/train/exp/weights/best.pt) model.export(formatonnx, opset12)导出后可以用onnxruntime验证输出节点数import onnxruntime as ort sess ort.InferenceSession(best.onnx) for inp in sess.get_inputs(): print(input:, inp.name, inp.shape) for out in sess.get_outputs(): print(output:, out.name, out.shape)如果修改成功这里应该能看到两个输出节点分别对应回归结果和分类结果。4.5 训练验证以 COCO128 这个小数据集为例快速验证改造后的流程能不能跑通yolo detect train datacoco128.yaml modelyolov8s.yaml epochs20 imgsz640如果数据集路径不对可以先下载 COCO128from ultralytics.data import utils utils.download_dataset(https://ultralytics.com/assets/coco128.zip)训练完成后重点检查两点训练 loss 是否正常下降。导出 ONNX 时输出节点是否为两个。如果 loss 曲线明显震荡或不下降优先检查是不是双头后的 loss 计算维度不匹配。5. RK3588 部署ONNX 导出与 RKNN 转换5.1 ONNX 导出注意事项在改完双头输出后ONNX 导出有几个坑要注意。第一个是 opset 版本。RKNN Toolkit2 对 ONNX opset 的兼容性有限一般建议固定在 12 或 13。太高版本可能导致某些算子无法解析。第二个是动态维度。因为 RKNN 转换时通常要指定固定输入尺寸建议导出时就把 batch 和宽高固定model.export(formatonnx, opset12, imgsz[640, 640], batch1)第三个是多输出节点的命名。导出后最好确认一下输出的名称因为后面 RKNN 转换时可能要用名称索引。5.2 RKNN Toolkit2 转换配置RK3588 对应的是 rknn-toolkit2注意不是老版本 rknn-toolkit。在 PC 上安装 rknn-toolkit2 后写一个转换脚本from rknn.api import RKNN rknn RKNN() # 配置量化与预处理的参数 rknn.config( mean_values[[0, 0, 0]], std_values[[255, 255, 255]], target_platformrk3588 ) # 加载 ONNX 模型 ret rknn.load_onnx(modelbest.onnx) if ret ! 0: print(load onnx failed) exit(1) # 构建 RKNN 模型 ret rknn.build(do_quantizationTrue, datasetdataset.txt) if ret ! 0: print(build failed) exit(1) # 导出 RKNN 模型 ret rknn.export_rknn(best.rknn) if ret ! 0: print(export failed) exit(1) rknn.release()其中dataset.txt是量化校准数据集列表每行一张图片路径img_001.jpg img_002.jpg img_003.jpg量化校准集的图片建议从训练集中挑选 200 到 500 张覆盖不同光照和目标形态。5.3 板端推理与后处理把best.rknn拷贝到 RK3588 板子上使用 RKNN Python API 做推理。核心代码如下import numpy as np from rknnlite.api import RKNNLite rknn RKNNLite() ret rknn.load_rknn(best.rknn) if ret ! 0: print(load rknn failed) exit(1) ret rknn.init_runtime() if ret ! 0: print(init runtime failed) exit(1) # 读取并预处理图片 img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img img.astype(np.float32) / 255.0 img img.transpose(2, 0, 1)[None] # 推理outputs 是一个列表 outputs rknn.inference(inputs[img]) # 第一个输出回归结果 # 第二个输出分类结果 reg_out outputs[0] cls_out outputs[1] print(reg_out shape:, reg_out.shape) print(cls_out shape:, cls_out.shape)这时候后处理就变得非常直观回归输出直接 reshape 成[1, 8400, 4]分类输出 reshape 成[1, 8400, nc]。然后做阈值过滤、NMS 就可以了。NMS 部分可以沿用 YOLOv8 官方后处理思路也可以集成一些端侧加速库。6. 常见问题与排查下面这些问题是双头改造和 RK3588 部署过程中比较常遇到的整理成一个排查表格。问题现象常见原因解决思路训练 loss 不下降双头输出后 loss 计算维度不匹配打印两个分支的 shape逐一核对导出 ONNX 后只有一个输出节点forward 里还是默认 cat 逻辑检查 Detect 类 forward 是否真的改到位ONNX 转 RKNN 失败opset 版本太高或不支持的算子把 opset 设为 12检查是否有自定义算子RKNN 推理输出 shape 不对输出节点索引和预期不一致用 Netron 打开 ONNX 查看节点顺序量化后精度下降明显校准图片太少或不具代表性增加校准集图片覆盖多种场景后处理出框位置偏移输出坐标解码方式变了确认 DFL 解码和 make_anchors 坐标逻辑6.1 训练 loss 不下降比较隐蔽的原因是make_anchors不同的调用方式导致 offset 不对。YOLOv8 官方在Detect的 forward 里使用make_anchors(x, self.stride, 0.5)如果你的双头实现里缺少这一步边界框坐标的解码结果就会乱掉。排查时可以在训练脚本里把预测结果打印出来看 DFL 解码后的坐标范围是否在合理区间内。如果坐标数值明显异常基本可以确定是 anchor 生成的问题。6.2 转 RKNN 报算子不支持YOLOv8 本身算子比较简单一般不会出现不支持的算子。但如果你的模型结构里加了自定义模块比如注意力机制或其他复杂操作转 RKNN 时可能报错。解决思路有两种一是简化模型结构把自定义算子用 RKNN 支持的常见算子重写。二是使用 RKNN 的算子自定义扩展功能。不过这个开发成本较高不建议在初期尝试。6.3 量化精度问题RK3588 上做 INT8 量化后精度下降在 1% 到 3% 以内通常是可接受的。如果下降太多优先做以下检查校准集是否和训练集来自同一分布。预处理是否一致。YOLOv8 训练时的归一化方式是除以 255RKNN 配置里的std_values要对应设置成 255。是否有对量化特别不友好的层比如输出的数值范围很大。可以尝试对部分层做混合精度量化。7. 工程建议与最佳实践7.1 结构设计建议双头输出不是越多越好。每增加一个输出头模型参数量和计算量都会增加对 RK3588 这种嵌入式平台来说推理耗时也会随之上升。在改造之前先明确业务痛点是多任务需求还是小目标需求是精度瓶颈还是后处理复杂度瓶颈明确需求后再确定双头的具体形态避免无意义的结构膨胀。7.2 训练技巧双头模型训练时两个分支的重要程度可能不一样。建议给分类分支和回归分支设置不同的 loss 权重防止一个分支的梯度过大压过另一个分支。例如在 loss 函数里loss_cls self.bce(cls_preds, cls_labels) loss_reg self.dfl(reg_preds, reg_labels) loss loss_cls * 0.5 loss_reg * 1.0具体的权重系数需要根据数据集特点调整。类别不平衡严重时分类 loss 的权重可以适当提高。另外双头模型对初始学习率比较敏感。建议在预训练权重的基础上做 fine-tune初始学习率不要太高默认的 0.01 通常没问题。7.3 部署侧建议在 RK3588 上部署双头模型时有几个工程细节值得注意。第一输入分辨率尽量固定。虽然 RKNN 支持动态输入但固定输入尺寸可以获得更好的 NPU 计算性能和更低的首帧延迟。第二多线程和零拷贝。如果做实时视频流检测建议用RKNNLite的多实例模式每个线程绑定一个核心配合 RK3588 的六核 CPU 和 NPU 并行性能上限会高很多。第三后处理尽量向量化。Python 环境下逐锚点遍历 8400 个结果会很慢建议用 NumPy 的向量化操作或者直接改写成 C 扩展。7.4 代码维护与版本管理这类对官方源码的修改强烈建议用 Git 管理并且把改动点集中在少数几个文件里。每次升级 ultralytics 版本时改动点越少合并成本越低。保存模型配置时把 yaml 文件和改过的 head.py 一起备份。RKNN 转换时模型结构依赖 yaml训练超参数依赖 yaml数据集配置依赖 yaml。没有 yaml后续复现和调试会很被动。8. 总结与下一步本文围绕 RK3588 YOLOv8 的端侧部署场景完整拆解了 YOLOv8 输出层的结构解释了单头输出和双头输出的区别并给出了从源码修改、训练、导出 ONNX、转 RKNN 到板端推理的完整流程。如果你也在做 RK3588 端侧目标检测并且有自定义输出结构的需求这套改造思路可以直接参考。所谓单头变双头本质上不是“两个检测框输出”而是把模型的任务空间拆得更细让后续的工程部署更加灵活。真正理解了Detect类的 forward 流程后你还可以举一反三改成第三个头、第四个头或者在输出层加入后处理需要的附加信息。下一步可以继续深入的方向在双头基础上加一个关键点输出头实现检测 姿态估计。把分类分支改成多标签分类用于同时输出多个属性。尝试将双头模型和 DeepSORT 等跟踪算法结合做实时多目标跟踪。如果在改动的过程中遇到问题欢迎在评论区交流。动手改一遍源码比看十篇教程都管用。 RK3588端侧部署YOLOv8：单头改双头输出实战指南
在做 RK3588 端侧部署时很多朋友会遇到一个很尴尬的情况模型在电脑上跑得好好的精度也够一旦转到 RKNN 就出现各种算子不支持、后处理写法别扭、多任务输出没法拆分的问题。尤其是 YOLOv8 这种自带解耦头的模型默认输出结构在边缘侧并不总是“最优解”。本篇文章就围绕 RK3588 YOLOv8 这个组合从源码层面拆解 YOLOv8 输出层的设计再一步步把默认的单头输出改造成双头输出并完整走一遍训练、导出、RKNN 转换和板端部署流程。这篇文章适合已经跑通 RK3588 基础环境、正在做模型优化或想自定义 YOLOv8 检测头的开发者。零基础的朋友可以先补一下 RKNN Toolkit2 的基础用法再回来看本文的改动思路。1. 为什么要拆解输出层从单头到双头的实际场景1.1 输出层在 YOLOv8 中扮演的角色YOLOv8 的网络结构可以粗分成三块Backbone主干网络、Neck特征融合、Head检测头。Backbone 负责提取图像特征Neck 把不同尺度的特征做融合最后交给 Head 输出目标类别和边界框位置。输出层也就是 Head 部分直接决定了“网络最终吐出什么”。默认情况下YOLOv8 的 Head 会在三个尺度上分别输出预测结果。这三个尺度对应特征图的下采样倍数通常是 8 倍、16 倍、32 倍。小幅度的特征图负责检测大目标大幅度负责检测小目标。在 YOLOv8 的源码里这个 Head 是由Detect类实现的。它的核心逻辑不复杂对每个尺度的特征图做卷积预测出目标的类别概率和边界框坐标最后把多尺度的结果拼到一起形成最终输出。如果只是做标准的目标检测直接用默认 Head 没有任何问题。但一旦涉及到“检测 分类 关键点”这样的多任务场景或者想把小目标检测单独拆出来做优化默认的单一输出结构就不够灵活了。1.2 单头与双头的本质区别先解释一下本文语境下的“单头”和“双头”。默认 YOLOv8 的Detect类虽然在内部把分类和回归分支分开了但从外部看最终只输出一个融合后的结果这个叫单头输出。它的结构大致是这样的输入特征图 ├── 分类分支Conv Conv Conv2d └── 回归分支Conv Conv Conv2d ↓ 拼接成一个 Tensor 输出所谓双头输出就是把这个单一结果拆成两个独立的输出分支。比如输入特征图 ├── 输出头 A负责类别预测 └── 输出头 B负责边界框回归还可以是另一种双头在原本的检测头基础上额外增加一个“辅助头”比如一个专门输出小目标检测结果的浅层头或者一个输出关键点坐标的头。在 RK3588 上做双头改造常见动机有三个多任务输出一个头做检测一个头做属性分类或者一个头做检测一个头做关键点。小目标增强在原有 P3/P4/P5 三个输出尺度的基础上增加一个 P2 浅层输出头让模型对小目标更敏感。NPU 部署优化把输出拆开后RKNN 可以分别处理不同尺度的结果后处理代码更容易编写和维护。1.3 RK3588 部署场景下的真实需求RK3588 内置的 NPU 算力虽然不错但和桌面级 GPU 不同它对网络结构的兼容性有自己的脾气。很多在 PyTorch 里能跑的算子转到 RKNN 时可能就不支持或者需要特殊写法。我在实际项目中遇到的一个典型场景是需要同时检测目标并判断目标朝向。原来的做法是单独训练一个 YOLOv8 检测模型和一个分类模型部署时串行推理。这样做有两个问题一是耗时翻倍二是两套模型的预处理和后处理互相独立工程上很繁琐。改成双头输出后检测和朝向分类在一个模型里完成前处理只做一次NPU 推理也只做一次后处理从两个模型的结果里各自取数据即可。整体帧率提升很明显。另一个常见场景是小目标检测。标准 YOLOv8 的三个输出尺度对密集小目标效果有限。很多做工厂缺陷检测、无人机航拍目标检测的朋友都会尝试加一个 P2 输出头。这个改动其实也可以理解成一种“双头到多头”的扩展思路。所以学会从源码层面修改 YOLOv8 的输出层是做端侧模型定制绕不开的基本功。2. 读懂 YOLOv8 的检测头源码2.1 YOLOv8 的 Decoupled Head 结构YOLOv8 和 YOLOv5 一个很大的区别就是 Head 从耦合头改成了解耦头。YOLOv5 的检测头是直接在特征图上预测(4 1 nc) * 3个通道其中 4 是边界框坐标1 是置信度nc 是类别数3 是 anchor 数量。这种方式把分类和回归耦合在同一个输出通道上。YOLOv8 则把分类和回归拆成两个独立分支。在源码ultralytics/nn/modules/head.py的Detect类中可以看到两个核心卷积模块cv2回归分支输出4 * reg_max个通道cv3分类分支输出nc个通道这里的reg_max是 YOLOv8 度量学习中 DFLDistribution Focal Loss模块的参数默认值是 16。也就是说每个边界框的每条边不是直接回归一个坐标值而是回归 16 个离散概率再通过 DFL 层加权求和得到最终坐标。这种结构图大致如下输入特征图P3/P4/P5 ├── cv2 分支Conv - Conv - Conv2d(4 * reg_max) │ ↓ │ 4 * reg_max 通道的回归特征 └── cv3 分支Conv - Conv - Conv2d(nc) ↓ nc 通道的分类特征2.2 Detect 类的初始化与 Forward 流程打开 YOLOv8 源码Detect 类的初始化函数中比较关键的部分如下class Detect(nn.Module): def __init__(self, nc80, ch()): super().__init__() self.nc nc self.nl len(ch) self.reg_max 16 self.no nc self.reg_max * 4 self.stride torch.zeros(self.nl) c2, c3 max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100)) self.cv2 nn.ModuleList( nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch ) self.cv3 nn.ModuleList( nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch ) self.dfl DFL(self.reg_max)这里ch是 Neck 输出的三个特征图的通道数。cv2和cv3都是 ModuleList每个特征图对应一套自己的卷积参数。forward方法做的事可以拆成下面几步遍历三个尺度的特征图分别经过cv2和cv3。对分类结果做 sigmoid 激活。对回归结果通过 DFL 计算边界框坐标。按尺度拼接并统一转成推理格式输出。核心代码如下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 ... # 推理模式拼接所有尺度的输出并返回注意这个torch.cat它把回归分支和分类分支拼接在通道维度上。这就是“单头”输出的核心所在外部调用模型时拿到的是一个已经拼好的结果。2.3 输出 Tensor 的形状与含义推理模式下YOLOv8 输出的格式是[batch, num_anchors, 4 nc]。num_anchors是所有尺度特征图的锚点总数。以 640x640 输入为例P3 特征图是 80x80有 6400 个锚点P4 特征图是 40x40有 1600 个锚点P5 特征图是 20x20有 400 个锚点合计 8400 个锚点。每个锚点对应4 nc个数值前 4 个是边界框坐标归一化的 x, y, w, h后面 nc 个是类别概率。了解这个形状很重要因为改造成双头输出后前 4 个数值和后面 nc 个数值会被拆到两个不同的 Tensor 中后处理代码需要相应调整。3. 环境准备与代码基线3.1 硬件与软件环境本文的环境如下版本可以根据自己的项目情况调整项目版本开发板RK3588本文以正点原子 RK3588 为例NPU 工具链rknn-toolkit2训练环境Python 3.8 PyTorch 2.0YOLOv8 源码ultralytics 8.x模型导出格式ONNX - RKNN操作系统可以是 Ubuntu也可以直接在 RK3588 的板端系统上做推理验证。训练和导出建议放在 PC 上完成RK3588 只负责 RKNN 模型推理。3.2 准备 YOLOv8 源码与预训练权重首先克隆 ultralytics 源码。如果网络不方便也可以直接在 PyPI 安装pip install ultralytics不过由于我们需要改源码里的Detect类建议直接使用 GitHub 源码方式git clone https://github.com/ultralytics/ultralytics.git cd ultralytics pip install -e .下载一个预训练权重作为基线比如 YOLOv8swget https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s.pt这个权重是 COCO 80 类预训练的后面做双头改造时我们可以先用它跑通流程再用自己的数据集重新训练。3.3 快速验证当前输出写一个最简单的推理脚本确认当前单头输出的形状import torch from ultralytics import YOLO model YOLO(yolov8s.pt) model.eval() dummy_input torch.randn(1, 3, 640, 640) with torch.no_grad(): result model.predict(dummy_input) print(type(result))如果要看更底层的模型输出形状可以直接拿model.model的 forward 结果import torch from ultralytics.nn.tasks import DetectionModel model DetectionModel(yolov8s.yaml) model.eval() dummy_input torch.randn(1, 3, 640, 640) with torch.no_grad(): output model(dummy_input) print(len(output)) for i, x in enumerate(output): print(i, x.shape)这里输出的就是推理阶段模型拼接后的最终结果。一共一个 Tensor形状是[1, 8400, 84]其中 84 4 80。记住这个基线结果后面改成双头输出后比较差异就有依据了。4. 核心改造单头变双头输出4.1 设计双头结构在动手改代码之前先把目标说清楚。我们这次要实现的“双头”是把默认的检测头从“一个融合输出”改造成“两个独立输出”输出头 A边界框回归结果形状为[batch, 8400, 4]输出头 B分类结果形状为[batch, 8400, nc]后处理时先从两个输出里分别取所需数据再合并计算最终检测框。这样做的好处是后处理代码更清晰不需要在一大块连续内存里切分。任务解耦后可以单独对分类分支做蒸馏、量化或优化。如果后续要扩展第三个头比如关键点头结构上也很方便。4.2 修改 Detect 类yolov8 的模型定义中检测头是动态加载的。我们需要修改ultralytics/nn/modules/head.py里的Detect类。核心思路是在forward中不再把cv2和cv3的结果cat在一起而是分别返回。下面给出修改后的核心代码class Detect(nn.Module): def __init__(self, nc80, ch()): super().__init__() self.nc nc self.nl len(ch) self.reg_max 16 self.no nc self.reg_max * 4 self.stride torch.zeros(self.nl) c2, c3 max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100)) self.cv2 nn.ModuleList( nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch ) self.cv3 nn.ModuleList( nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch ) self.dfl DFL(self.reg_max) def forward(self, x): shape x[0].shape dfl_out [] cls_out [] for i in range(self.nl): reg_branch self.cv2[i](x[i]) cls_branch self.cv3[i](x[i]) # 回归分支解码 b, c, h, w reg_branch.shape reg_branch reg_branch.view(b, 4, self.reg_max, h * w).transpose(2, 3).reshape(b, 4, h * w) dfl_out.append(self.dfl(reg_branch)) # 分类分支 sigmoid cls_out.append(torch.sigmoid(cls_branch.reshape(b, self.nc, h * w))) if self.training: # 训练阶段返回回归原始特征和分类特征便于 loss 计算 return dfl_out, cls_out # 推理阶段分尺度转换为 xywh 格式 anchors, strides self.make_anchors(x, self.stride, 0.5) results [] for i in range(self.nl): d dfl_out[i].transpose(1, 2) c cls_out[i].transpose(1, 2) results.append((d, c)) return results这段代码里有几个关键点需要解释一下reg_branch要经过 DFL 解码才能得到边界框坐标。cls_branch需要做 sigmoid得到类别概率。训练阶段和推理阶段返回的格式不一样。训练阶段需要保留原始特征方便 loss 计算推理阶段输出解码后的坐标和类别概率方便后处理直接使用。当然实际工程中很多人不会直接在主源码里大改而是继承Detect写一个DetectDual类。这样改动更小也不容易影响原有逻辑。比如class DetectDual(Detect): def forward(self, x): # 重写 forward拆分输出 ...然后在模型的 yaml 配置里把 head 的类型改成DetectDual。如果你用的是自定义数据集还要注意nc改成自己的类别数。比如检测车辆和行人nc2。4.3 修改 Loss 计算逻辑输出结构变了loss 计算也必须跟着改。默认情况下YOLOv8 的 loss 变化在ultralytics/utils/loss.py文件里调用。原来的forward返回拼接结果loss 会从拼接结果里切分分类分支和回归分支。改成双头输出后loss 函数需要直接接收两个分支的结果。修改原则是分类 loss 只依赖分类分支输出回归 loss 只依赖回归分支输出。点开 v8DetectionLoss 类把原先从pred中切分pred_distri和pred_scores的部分改成直接接收双头结果。大致逻辑如下class v8DetectionLoss: def __init__(self, model): # 原有初始化逻辑保持不变 ... def __call__(self, preds, batch): # preds 变成了 (dfl_out, cls_out) pred_distri, pred_scores preds ...需要注意的是训练阶段我们返回的是未解码的回归特征loss 内部会通过 DFL 计算分布损失。保持和官方一致的流程即可。4.4 修改导出逻辑模型训练完以后要导出 ONNX。导出的逻辑在ultralytics/nn/tasks.py的Model.export或exporter.py里。默认导出时会把 forward 的输出拼接成一个维度输出。改双头后导出时会输出两个节点这符合我们的需求。导出命令yolo export modelruns/train/exp/weights/best.pt formatonnx opset12或者直接用 Python APIfrom ultralytics import YOLO model YOLO(runs/train/exp/weights/best.pt) model.export(formatonnx, opset12)导出后可以用onnxruntime验证输出节点数import onnxruntime as ort sess ort.InferenceSession(best.onnx) for inp in sess.get_inputs(): print(input:, inp.name, inp.shape) for out in sess.get_outputs(): print(output:, out.name, out.shape)如果修改成功这里应该能看到两个输出节点分别对应回归结果和分类结果。4.5 训练验证以 COCO128 这个小数据集为例快速验证改造后的流程能不能跑通yolo detect train datacoco128.yaml modelyolov8s.yaml epochs20 imgsz640如果数据集路径不对可以先下载 COCO128from ultralytics.data import utils utils.download_dataset(https://ultralytics.com/assets/coco128.zip)训练完成后重点检查两点训练 loss 是否正常下降。导出 ONNX 时输出节点是否为两个。如果 loss 曲线明显震荡或不下降优先检查是不是双头后的 loss 计算维度不匹配。5. RK3588 部署ONNX 导出与 RKNN 转换5.1 ONNX 导出注意事项在改完双头输出后ONNX 导出有几个坑要注意。第一个是 opset 版本。RKNN Toolkit2 对 ONNX opset 的兼容性有限一般建议固定在 12 或 13。太高版本可能导致某些算子无法解析。第二个是动态维度。因为 RKNN 转换时通常要指定固定输入尺寸建议导出时就把 batch 和宽高固定model.export(formatonnx, opset12, imgsz[640, 640], batch1)第三个是多输出节点的命名。导出后最好确认一下输出的名称因为后面 RKNN 转换时可能要用名称索引。5.2 RKNN Toolkit2 转换配置RK3588 对应的是 rknn-toolkit2注意不是老版本 rknn-toolkit。在 PC 上安装 rknn-toolkit2 后写一个转换脚本from rknn.api import RKNN rknn RKNN() # 配置量化与预处理的参数 rknn.config( mean_values[[0, 0, 0]], std_values[[255, 255, 255]], target_platformrk3588 ) # 加载 ONNX 模型 ret rknn.load_onnx(modelbest.onnx) if ret ! 0: print(load onnx failed) exit(1) # 构建 RKNN 模型 ret rknn.build(do_quantizationTrue, datasetdataset.txt) if ret ! 0: print(build failed) exit(1) # 导出 RKNN 模型 ret rknn.export_rknn(best.rknn) if ret ! 0: print(export failed) exit(1) rknn.release()其中dataset.txt是量化校准数据集列表每行一张图片路径img_001.jpg img_002.jpg img_003.jpg量化校准集的图片建议从训练集中挑选 200 到 500 张覆盖不同光照和目标形态。5.3 板端推理与后处理把best.rknn拷贝到 RK3588 板子上使用 RKNN Python API 做推理。核心代码如下import numpy as np from rknnlite.api import RKNNLite rknn RKNNLite() ret rknn.load_rknn(best.rknn) if ret ! 0: print(load rknn failed) exit(1) ret rknn.init_runtime() if ret ! 0: print(init runtime failed) exit(1) # 读取并预处理图片 img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (640, 640)) img img.astype(np.float32) / 255.0 img img.transpose(2, 0, 1)[None] # 推理outputs 是一个列表 outputs rknn.inference(inputs[img]) # 第一个输出回归结果 # 第二个输出分类结果 reg_out outputs[0] cls_out outputs[1] print(reg_out shape:, reg_out.shape) print(cls_out shape:, cls_out.shape)这时候后处理就变得非常直观回归输出直接 reshape 成[1, 8400, 4]分类输出 reshape 成[1, 8400, nc]。然后做阈值过滤、NMS 就可以了。NMS 部分可以沿用 YOLOv8 官方后处理思路也可以集成一些端侧加速库。6. 常见问题与排查下面这些问题是双头改造和 RK3588 部署过程中比较常遇到的整理成一个排查表格。问题现象常见原因解决思路训练 loss 不下降双头输出后 loss 计算维度不匹配打印两个分支的 shape逐一核对导出 ONNX 后只有一个输出节点forward 里还是默认 cat 逻辑检查 Detect 类 forward 是否真的改到位ONNX 转 RKNN 失败opset 版本太高或不支持的算子把 opset 设为 12检查是否有自定义算子RKNN 推理输出 shape 不对输出节点索引和预期不一致用 Netron 打开 ONNX 查看节点顺序量化后精度下降明显校准图片太少或不具代表性增加校准集图片覆盖多种场景后处理出框位置偏移输出坐标解码方式变了确认 DFL 解码和 make_anchors 坐标逻辑6.1 训练 loss 不下降比较隐蔽的原因是make_anchors不同的调用方式导致 offset 不对。YOLOv8 官方在Detect的 forward 里使用make_anchors(x, self.stride, 0.5)如果你的双头实现里缺少这一步边界框坐标的解码结果就会乱掉。排查时可以在训练脚本里把预测结果打印出来看 DFL 解码后的坐标范围是否在合理区间内。如果坐标数值明显异常基本可以确定是 anchor 生成的问题。6.2 转 RKNN 报算子不支持YOLOv8 本身算子比较简单一般不会出现不支持的算子。但如果你的模型结构里加了自定义模块比如注意力机制或其他复杂操作转 RKNN 时可能报错。解决思路有两种一是简化模型结构把自定义算子用 RKNN 支持的常见算子重写。二是使用 RKNN 的算子自定义扩展功能。不过这个开发成本较高不建议在初期尝试。6.3 量化精度问题RK3588 上做 INT8 量化后精度下降在 1% 到 3% 以内通常是可接受的。如果下降太多优先做以下检查校准集是否和训练集来自同一分布。预处理是否一致。YOLOv8 训练时的归一化方式是除以 255RKNN 配置里的std_values要对应设置成 255。是否有对量化特别不友好的层比如输出的数值范围很大。可以尝试对部分层做混合精度量化。7. 工程建议与最佳实践7.1 结构设计建议双头输出不是越多越好。每增加一个输出头模型参数量和计算量都会增加对 RK3588 这种嵌入式平台来说推理耗时也会随之上升。在改造之前先明确业务痛点是多任务需求还是小目标需求是精度瓶颈还是后处理复杂度瓶颈明确需求后再确定双头的具体形态避免无意义的结构膨胀。7.2 训练技巧双头模型训练时两个分支的重要程度可能不一样。建议给分类分支和回归分支设置不同的 loss 权重防止一个分支的梯度过大压过另一个分支。例如在 loss 函数里loss_cls self.bce(cls_preds, cls_labels) loss_reg self.dfl(reg_preds, reg_labels) loss loss_cls * 0.5 loss_reg * 1.0具体的权重系数需要根据数据集特点调整。类别不平衡严重时分类 loss 的权重可以适当提高。另外双头模型对初始学习率比较敏感。建议在预训练权重的基础上做 fine-tune初始学习率不要太高默认的 0.01 通常没问题。7.3 部署侧建议在 RK3588 上部署双头模型时有几个工程细节值得注意。第一输入分辨率尽量固定。虽然 RKNN 支持动态输入但固定输入尺寸可以获得更好的 NPU 计算性能和更低的首帧延迟。第二多线程和零拷贝。如果做实时视频流检测建议用RKNNLite的多实例模式每个线程绑定一个核心配合 RK3588 的六核 CPU 和 NPU 并行性能上限会高很多。第三后处理尽量向量化。Python 环境下逐锚点遍历 8400 个结果会很慢建议用 NumPy 的向量化操作或者直接改写成 C 扩展。7.4 代码维护与版本管理这类对官方源码的修改强烈建议用 Git 管理并且把改动点集中在少数几个文件里。每次升级 ultralytics 版本时改动点越少合并成本越低。保存模型配置时把 yaml 文件和改过的 head.py 一起备份。RKNN 转换时模型结构依赖 yaml训练超参数依赖 yaml数据集配置依赖 yaml。没有 yaml后续复现和调试会很被动。8. 总结与下一步本文围绕 RK3588 YOLOv8 的端侧部署场景完整拆解了 YOLOv8 输出层的结构解释了单头输出和双头输出的区别并给出了从源码修改、训练、导出 ONNX、转 RKNN 到板端推理的完整流程。如果你也在做 RK3588 端侧目标检测并且有自定义输出结构的需求这套改造思路可以直接参考。所谓单头变双头本质上不是“两个检测框输出”而是把模型的任务空间拆得更细让后续的工程部署更加灵活。真正理解了Detect类的 forward 流程后你还可以举一反三改成第三个头、第四个头或者在输出层加入后处理需要的附加信息。下一步可以继续深入的方向在双头基础上加一个关键点输出头实现检测 姿态估计。把分类分支改成多标签分类用于同时输出多个属性。尝试将双头模型和 DeepSORT 等跟踪算法结合做实时多目标跟踪。如果在改动的过程中遇到问题欢迎在评论区交流。动手改一遍源码比看十篇教程都管用。