1. 背景与痛点
在 Rockchip NPU(如 RK3588)上部署 YOLO 模型时,官方推荐的做法是使用瑞芯微优化过的ultralytics_yolo11仓库导出 ONNX。但这种方式存在几个明显的痛点:
- 仓库依赖:必须克隆整个
ultralytics_yolo11仓库,不能使用标准的pip install ultralytics。 - 环境隔离:需要单独维护一个 conda 环境,且容易与其他项目产生冲突。
- 版本锁定:fork 仓库的更新节奏与官方不同,可能滞后于上游的 bug 修复和新特性。
- 调试困难:修改源码后想复用到另一个项目,往往需要重新处理依赖关系。
本文将分享一个独立脚本,它基于官方ultralytics库,通过monkey‑patch动态注入 RKNN 优化逻辑,无需依赖任何 fork 仓库。一条命令,就能将你的 YOLOv11 模型导出为适配 Rockchip NPU 的 ONNX 格式。
2. 脚本核心原理
2.1 优化目标
根据瑞芯微官方的优化说明,导出的 ONNX 需要实现三项关键优化:
| 优化项 | 说明 |
|---|---|
| 移除后处理 | 把 NMS、Anchor 解码等留在 CPU 端(由rknn_model_zoo负责),推理侧只保留纯网络计算。 |
| 移除 DFL | DFL 结构对 NPU 不友好,改为直接输出原始的 box 特征图,减少算子适配压力。 |
| 新增 cls_sum 分支 | 输出每个位置类别概率的逐像素求和,用于加速 CPU 后处理阶段的阈值筛选。 |
2.2 实现方式
脚本的核心是重写检测头(Detect/Pose)的forward方法。以 Detect 头为例,修改后的逻辑如下:
def_detect_forward_rknn(self,x):"""返回 3 层 × 3 路输出 = 9 个 tensor"""y=[]foriinrange(self.nl):box=self.cv2[i](x[i])# (B, 4*reg_max, H, W)cls=torch.sigmoid(self.cv3[i](x[i]))# (B, nc, H, W)cls_sum=torch.clamp(cls.sum(1,keepdim=True),0,1)# (B, 1, H, W)y.extend([box,cls,cls_sum])returny对于 Pose(姿态估计)任务,额外处理了关键点解码:
def_pose_forward_rknn(self,x):# 1) 检测部分:走上面的 Detect 分支det_outs=_detect_forward_rknn(self,x,"Pose")# 2) 关键点:anchor 解码 + stride 缩放anchors,strides=make_anchors(det_outs,self.stride,0.5)kpt=self.cv4[i](x[i])# 原始关键点pred_kpt=(kpt*2.0+anchors-0.5)*stridesreturn[det_outs,pred_kpt]这两种改写保证了导出后的 ONNX 不再包含对 NPU 不友好的结构,同时让 CPU 端的后处理更加高效。
2.3 与官方仓库方式的对比
| 对比项 | 官方ultralytics_yolo11 | 本独立脚本 |
|---|---|---|
| 安装方式 | git clone+pip install -e . | pip install ultralytics |
| 依赖仓库 | 必须使用 airockchip fork | 任意官方ultralytics版本 |
| 灵活性 | 需要在 fork 仓库中修改源码 | 脚本独立,逻辑清晰,随时可改 |
| 复用性 | 绑定在特定目录结构下 | 单文件即可复制到任何项目中使用 |
3. 使用方法
3.1 环境准备
# 创建虚拟环境(可选但推荐)conda create-nyolo_exportpython=3.9-yconda activate yolo_export# 安装依赖pipinstallultralytics onnx onnxscript onnxsim3.2 脚本配置
修改脚本开头几行配置即可适配你的模型:
PT_PATH="best.pt"# 你的模型路径ONNX_PATH="best.onnx"# 输出路径IMG_SIZE=320# 训练时的输入尺寸OPSET=13# ONNX opset 版本(推荐 12~13)3.3 运行导出
python export_rknn_onnx.py执行后,脚本会在指定路径生成一个对 RKNN 友好的 ONNX 文件。
3.4 输出说明
- Pose 任务输出 4 个 tensor:
| 输出名 | 形状 | 说明 |
|---|---|---|
| det_p3 | (B, 4*reg_max+nc, 80, 80) | P3 层检测特征图 |
| det_p4 | (B, 4*reg_max+nc, 40, 40) | P4 层检测特征图 |
| det_p5 | (B, 4*reg_max+nc, 20, 20) | P5 层检测特征图 |
| kpts | (B, 17, 3, N) | 解码后的关键点 (x, y, visibility) |
- Detect 任务输出 9 个 tensor(3 层 × box/cls/cls_sum),此处不再赘述。
4. 验证 ONNX 模型
脚本会调用onnx.checker验证模型结构,并输出详细信息,便于快速核对输入输出形状:
📦 Opset 版本: domain: ai.onnx version: 13 IR version : 9 producer : pytorch 2.11.0 📥 输入: images shape=[1, 3, 320, 320] dtype=float32 📤 输出: det_p3 shape=[1, 69, 80, 80] dtype=float32 det_p4 shape=[1, 69, 40, 40] dtype=float32 det_p5 shape=[1, 69, 20, 20] dtype=float32 kpts shape=[1, 17, 3, 8400] dtype=float325. 后续部署
生成 ONNX 后,使用瑞芯微官方工具链即可完成后续部署:
# 1. 克隆 rknn_model_zoogitclone https://github.com/airockchip/rknn_model_zoo.git# 2. 转换为 RKNNcdrknn_model_zoo/examples/yolo11/python python convert.py--modelbest0531.onnx--targetrk3588--outputbest0531.rknn# 3. 推理测试python yolo11.py--modelbest0531.rknn--imagetest.jpg⚠️注意:由于模型已移除了后处理,推理后必须配合rknn_model_zoo提供的 CPU 后处理代码(即postprocess函数)才能得到最终的检测框与姿态点。
6. 完整脚本代码
""" 将 best.pt (YOLO11s-pose, 17 keypoints, 320x320) 导出为 RKNN 友好的 ONNX。 对照 airockchip 在 ultralytics 上的 RKNN 改造 (RKOPT_README.md): 1) 移除模型内的后处理(DFL / anchor decode / NMS),后处理放到 CPU 端做(rknn_model_zoo)。 2) 检测分支输出 3 个 stride 层的原始特征图: [cv2(x), sigmoid(cv3(x)), sum(sigmoid(cv3(x)))] - cv2 : box 分支输出 (B, 4*reg_max, H, W),不做 DFL - cv3 : cls 分支输出经过 sigmoid 后的逐类概率 (B, nc, H, W) - 第三个分支是 cls 概率的逐像素求和 (B, 1, H, W),用于后处理阶段加速阈值筛选 3) Pose 任务额外输出 17 个关键点的解码结果 (anchor + stride),(B, 17, 3, N) 为了脱离 ultralytics_yolo11-main 目录,本脚本不依赖该 fork —— 它从普通 pip 安装的 `ultralytics` 加载 .pt,然后对其检测头做 monkey-patch,以注入上述 RKNN 分支。 """importosimportsysimporttypesimporttorchimporttorch.nnasnn# ==================== 配置 ====================PT_PATH="best.pt"ONNX_PATH="best.onnx"IMG_SIZE=320# 与 trainYolo.py 一致OPSET=13DEVICE="cpu"# ==================== Detect / Pose 的 RKNN 前向 ====================def_detect_forward_rknn(self,x,task_type="Detect"):""" 照抄 airockchip ultralytics fork 中 head.Detect.forward 的 RKNN 分支。 返回结构(按 stride 层 P3/P4/P5 顺序排好): Pose / OBB 任务: [(cv2+cv3) cat raw, ...] per layer 共 3 个 tensor (B, 4*reg_max+nc, H, W) Detect 任务: [cv2_P3, sigmoid(cv3_P3), sum_cls_P3, cv2_P4, sigmoid(cv3_P4), sum_cls_P4, cv2_P5, sigmoid(cv3_P5), sum_cls_P5] """# Pose / OBB 走这一支:返回 cv2+cv3 拼接后的原始特征图(不做 DFL/sigmoid),# 关键点解码放在 Pose.forward 里iftask_typein("Pose","Obb"):y=[]foriinrange(self.nl):y.append(torch.cat((self.cv2[i](x[i]),self.cv3[i](x[i])),1))returny# 纯 Detect 分支: box / sigmoid(cls) / cls_sum 三件套y=[]foriinrange(self.nl):box=self.cv2[i](x[i])# (B, 4*reg_max, H, W)cls=torch.sigmoid(self.cv3[i](x[i]))# (B, nc, H, W)cls_sum=torch.clamp(cls.sum(1,keepdim=True),0,1)# (B, 1, H, W)y.append(box)y.append(cls)y.append(cls_sum)returnydef_pose_forward_rknn(self,x):""" 照抄 airockchip ultralytics fork 中 head.Pose.forward 的 RKNN 分支。 返回 [output_x, pred_kpt] output_x : list of 3 个 (B, 4*reg_max + nc, Hi, Wi) 原始检测特征图 pred_kpt : (B, num_kpt, 3, N) anchor 解码后的关键点 通道顺序 [x_decoded, y_decoded, sigmoid(visibility)] 其中 N = sum(Hi * Wi) 即所有 stride 层格点数之和 """bs=x[0].shape[0]# 关键点原始输出 (B, nk, N)kpt=torch.cat([self.cv4[i](x[i]).view(bs,self.nk,-1)foriinrange(self.nl)],dim=-1,)# 1) 检测部分:走 task_type='Pose' 分支 -> 返回 3 个原始 cv2+cv3 特征图(不做 DFL)output_x=_detect_forward_rknn(self,x,"Pose")# 2) 构造 anchors / strides(不依赖 self._inference,避免不同 ultralytics 版本 API 差异)fromultralytics.utils.talimportmake_anchors anchors,strides=make_anchors(output_x,self.stride,0.5)# anchors: (N, 2) strides: (N, 1) N = sum(Hi*Wi)# 转成 (2, N) / (1, N) 以匹配 kpts_decode 的逻辑anchors=anchors.transpose(0,1).contiguous()# (2, N)strides=strides.transpose(0,1).contiguous()# (1, N)# 3) 关键点解码:复刻 kpts_decode 的 RKNN 分支ndim=self.kpt_shape[1]y=kpt.view(bs,*self.kpt_shape,-1)# (B, num_kpt, ndim, N)a=(y[:,:,:2]*2.0+(anchors-0.5))*strides# (B, num_kpt, 2, N)ifndim==3:a=torch.cat((a,y[:,:,2:3].sigmoid()),2)# (B, num_kpt, 3, N)pred_kpt=areturn[output_x,pred_kpt]def_patch_head_for_rknn(model):""" 把模型最后一个检测头(Detect/Pose/Segment/OBB)切到 RKNN 导出模式。 仅依赖 ultralytics 官方包的 head 模块类名,不依赖那个 fork 的源码。 """fromultralytics.nn.modules.headimportDetect,Pose,Segment,OBB head=Noneforminmodel.modules():ifisinstance(m,(Detect,Pose,Segment,OBB)):head=m# 取最后一个,YOLO 通常只有一个检测头ifheadisNone:raiseRuntimeError("未在模型中找到 Detect/Pose/Segment/OBB 检测头")head.export=Truehead.format="rknn"head.dynamic=False# monkey-patch forward(只换 Detect / Pose,本仓库训练的是 pose)head.forward=types.MethodType(_pose_forward_rknnifisinstance(head,Pose)else_detect_forward_rknn,head)print(f"✅ 已为{type(head).__name__}头注入 RKNN 导出分支 "f"(nc={getattr(head,'nc','?')}, reg_max={getattr(head,'reg_max','?')}, "f"kpt_shape={getattr(head,'kpt_shape','?')})")returnhead# ==================== 导出主流程 ====================defexport(pt_path=PT_PATH,onnx_path=ONNX_PATH,imgsz=IMG_SIZE,opset=OPSET,device=DEVICE):print(f"🚀 开始导出 RKNN 友好的 ONNX")print(f" 权重 :{pt_path}")print(f" 输出 :{onnx_path}")print(f" 输入尺寸 :{imgsz}x{imgsz}")print(f" opset :{opset}")ifnotos.path.exists(pt_path):raiseFileNotFoundError(f"找不到权重文件:{pt_path}")# 用 ultralytics 官方 API 加载 .ptfromultralyticsimportYOLO yolo=YOLO(pt_path)model=yolo.model.to(device).eval().float()# 关闭 train 状态、推理用的 graph,但保留 BN/Conv 的融合ifhasattr(model,"fuse"):try:model=model.fuse()exceptExceptionase:print(f"⚠️ model.fuse() 失败,继续:{e}")forpinmodel.parameters():p.requires_grad=False# 注入 RKNN 分支head=_patch_head_for_rknn(model)# 干跑一次,确认输出结构 & 让 anchors 提前算好im=torch.zeros(1,3,imgsz,imgsz,dtype=torch.float32,device=device)withtorch.no_grad():out=model(im)# 打印输出结构print("\n📤 导出前 forward() 输出结构:")ifisinstance(out,(list,tuple))andlen(out)>0andisinstance(out[0],(list,tuple)):det_outs,kpt_out=out[0],out[1]fori,tinenumerate(det_outs):print(f" det[{i}] shape={tuple(t.shape)}")print(f" kpt shape={tuple(kpt_out.shape)}")else:fori,tinenumerate(out):print(f" out[{i}] shape={tuple(t.shape)}")# ===== 构造输出名 =====# Pose: [ [det_p3, det_p4, det_p5], pred_kpt ]# 展平为单层 list 才好命名class_PoseWrapper(nn.Module):def__init__(self,inner):super().__init__()self.inner=innerdefforward(self,x):det_outs,pred_kpt=self.inner(x)returndet_outs[0],det_outs[1],det_outs[2],pred_kptfromultralytics.nn.modules.headimportPoseifisinstance(head,Pose):wrapped=_PoseWrapper(model)output_names=["det_p3","det_p4","det_p5","kpts"]else:wrapped=model# 9 路输出(3 层 × box/cls/cls_sum)output_names=[]forsin("p3","p4","p5"):output_names+=[f"box_{s}",f"cls_{s}",f"cls_sum_{s}"]wrapped.eval()print(f"\n🛠️ torch.onnx.export →{onnx_path}")torch.onnx.export(wrapped,im,onnx_path,verbose=False,opset_version=opset,do_constant_folding=True,input_names=["images"],output_names=output_names,)print(f"✅ ONNX 已保存:{onnx_path}")_verify_onnx(onnx_path)print("\n📌 后续:把 .onnx 喂给 rknn-toolkit2 / rknn_model_zoo 即可转 .rknn")def_verify_onnx(onnx_path):try:importonnxfromonnximportTensorProtoexceptImportError:print("⚠️ 未安装 onnx 库,跳过模型验证 (`pip install onnx`)")returnonnx_model=onnx.load(onnx_path)onnx.checker.check_model(onnx_model)print("✅ ONNX 模型结构验证通过")# ---- opset 版本 ----print("\n📦 Opset 版本:")foropsetinonnx_model.opset_import:domain=opset.domainifopset.domainelse"ai.onnx"print(f" domain:{domain:20s}version:{opset.version}")# IR 版本与生产者print(f" IR version :{onnx_model.ir_version}")print(f" producer :{onnx_model.producer_name}{onnx_model.producer_version}")def_dtype_name(t):try:returnTensorProto.DataType.Name(t)exceptException:returnstr(t)def_shape(node):s=[]fordinnode.type.tensor_type.shape.dim:s.append(d.dim_valueifd.dim_valueelsed.dim_param)returnsprint("\n📥 输入:")forinpinonnx_model.graph.input:print(f"{inp.name:12s}shape={_shape(inp)}dtype={_dtype_name(inp.type.tensor_type.elem_type)}")print("\n📤 输出:")foroutinonnx_model.graph.output:print(f"{out.name:12s}shape={_shape(out)}dtype={_dtype_name(out.type.tensor_type.elem_type)}")if__name__=="__main__":export()📎 上述代码可直接复制保存为export_rknn_onnx.py使用。如果遇到具体的模型结构差异(如不同 YOLO 版本检测头命名),可在_patch_head_for_rknn附近稍作调整。
7. 总结
本文介绍的独立脚本,让你彻底告别对airockchip/ultralytics_yolo11仓库的依赖,仅需pip install ultralytics外加一个 Python 文件,即可生成对 RKNN 友好的 ONNX 模型。无论是快速验证 YOLO11 在 RK3588 上的效果,还是集成到自动化模型转换流水线中,这个方案都能显著降低成本、提升效率。欢迎直接拿去使用,遇到问题可以在项目仓库中反馈讨论。