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负责推理侧只保留纯网络计算。移除 DFLDFL 结构对 NPU 不友好改为直接输出原始的 box 特征图减少算子适配压力。新增 cls_sum 分支输出每个位置类别概率的逐像素求和用于加速 CPU 后处理阶段的阈值筛选。2.2 实现方式脚本的核心是重写检测头Detect/Pose的forward方法。以 Detect 头为例修改后的逻辑如下def_detect_forward_rknn(self,x):返回 3 层 × 3 路输出 9 个 tensory[]foriinrange(self.nl):boxself.cv2[i](x[i])# (B, 4*reg_max, H, W)clstorch.sigmoid(self.cv3[i](x[i]))# (B, nc, H, W)cls_sumtorch.clamp(cls.sum(1,keepdimTrue),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,stridesmake_anchors(det_outs,self.stride,0.5)kptself.cv4[i](x[i])# 原始关键点pred_kpt(kpt*2.0anchors-0.5)*stridesreturn[det_outs,pred_kpt]这两种改写保证了导出后的 ONNX 不再包含对 NPU 不友好的结构同时让 CPU 端的后处理更加高效。2.3 与官方仓库方式的对比对比项官方ultralytics_yolo11本独立脚本安装方式git clonepip install -e .pip install ultralytics依赖仓库必须使用 airockchip fork任意官方ultralytics版本灵活性需要在 fork 仓库中修改源码脚本独立逻辑清晰随时可改复用性绑定在特定目录结构下单文件即可复制到任何项目中使用3. 使用方法3.1 环境准备# 创建虚拟环境可选但推荐conda create-nyolo_exportpython3.9-yconda activate yolo_export# 安装依赖pipinstallultralytics onnx onnxscript onnxsim3.2 脚本配置修改脚本开头几行配置即可适配你的模型PT_PATHbest.pt# 你的模型路径ONNX_PATHbest.onnx# 输出路径IMG_SIZE320# 训练时的输入尺寸OPSET13# ONNX opset 版本推荐 12~133.3 运行导出python export_rknn_onnx.py执行后脚本会在指定路径生成一个对 RKNN 友好的 ONNX 文件。3.4 输出说明Pose 任务输出 4 个 tensor输出名形状说明det_p3(B, 4*reg_maxnc, 80, 80)P3 层检测特征图det_p4(B, 4*reg_maxnc, 40, 40)P4 层检测特征图det_p5(B, 4*reg_maxnc, 20, 20)P5 层检测特征图kpts(B, 17, 3, N)解码后的关键点 (x, y, visibility)Detect 任务输出 9 个 tensor3 层 × 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] dtypefloat32 输出: det_p3 shape[1, 69, 80, 80] dtypefloat32 det_p4 shape[1, 69, 40, 40] dtypefloat32 det_p5 shape[1, 69, 20, 20] dtypefloat32 kpts shape[1, 17, 3, 8400] dtypefloat325. 后续部署生成 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_PATHbest.ptONNX_PATHbest.onnxIMG_SIZE320# 与 trainYolo.py 一致OPSET13DEVICEcpu# Detect / Pose 的 RKNN 前向 def_detect_forward_rknn(self,x,task_typeDetect): 照抄 airockchip ultralytics fork 中 head.Detect.forward 的 RKNN 分支。 返回结构(按 stride 层 P3/P4/P5 顺序排好): Pose / OBB 任务: [(cv2cv3) cat raw, ...] per layer 共 3 个 tensor (B, 4*reg_maxnc, 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 走这一支:返回 cv2cv3 拼接后的原始特征图(不做 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):boxself.cv2[i](x[i])# (B, 4*reg_max, H, W)clstorch.sigmoid(self.cv3[i](x[i]))# (B, nc, H, W)cls_sumtorch.clamp(cls.sum(1,keepdimTrue),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 层格点数之和 bsx[0].shape[0]# 关键点原始输出 (B, nk, N)kpttorch.cat([self.cv4[i](x[i]).view(bs,self.nk,-1)foriinrange(self.nl)],dim-1,)# 1) 检测部分:走 task_typePose 分支 - 返回 3 个原始 cv2cv3 特征图(不做 DFL)output_x_detect_forward_rknn(self,x,Pose)# 2) 构造 anchors / strides(不依赖 self._inference,避免不同 ultralytics 版本 API 差异)fromultralytics.utils.talimportmake_anchors anchors,stridesmake_anchors(output_x,self.stride,0.5)# anchors: (N, 2) strides: (N, 1) N sum(Hi*Wi)# 转成 (2, N) / (1, N) 以匹配 kpts_decode 的逻辑anchorsanchors.transpose(0,1).contiguous()# (2, N)stridesstrides.transpose(0,1).contiguous()# (1, N)# 3) 关键点解码:复刻 kpts_decode 的 RKNN 分支ndimself.kpt_shape[1]ykpt.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)ifndim3:atorch.cat((a,y[:,:,2:3].sigmoid()),2)# (B, num_kpt, 3, N)pred_kptareturn[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 headNoneforminmodel.modules():ifisinstance(m,(Detect,Pose,Segment,OBB)):headm# 取最后一个,YOLO 通常只有一个检测头ifheadisNone:raiseRuntimeError(未在模型中找到 Detect/Pose/Segment/OBB 检测头)head.exportTruehead.formatrknnhead.dynamicFalse# monkey-patch forward(只换 Detect / Pose,本仓库训练的是 pose)head.forwardtypes.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,?)}, fkpt_shape{getattr(head,kpt_shape,?)}))returnhead# 导出主流程 defexport(pt_pathPT_PATH,onnx_pathONNX_PATH,imgszIMG_SIZE,opsetOPSET,deviceDEVICE):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 yoloYOLO(pt_path)modelyolo.model.to(device).eval().float()# 关闭 train 状态、推理用的 graph,但保留 BN/Conv 的融合ifhasattr(model,fuse):try:modelmodel.fuse()exceptExceptionase:print(f⚠️ model.fuse() 失败,继续:{e})forpinmodel.parameters():p.requires_gradFalse# 注入 RKNN 分支head_patch_head_for_rknn(model)# 干跑一次,确认输出结构 让 anchors 提前算好imtorch.zeros(1,3,imgsz,imgsz,dtypetorch.float32,devicedevice)withtorch.no_grad():outmodel(im)# 打印输出结构print(\n 导出前 forward() 输出结构:)ifisinstance(out,(list,tuple))andlen(out)0andisinstance(out[0],(list,tuple)):det_outs,kpt_outout[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.innerinnerdefforward(self,x):det_outs,pred_kptself.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:wrappedmodel# 9 路输出(3 层 × box/cls/cls_sum)output_names[]forsin(p3,p4,p5):output_names[fbox_{s},fcls_{s},fcls_sum_{s}]wrapped.eval()print(f\n️ torch.onnx.export →{onnx_path})torch.onnx.export(wrapped,im,onnx_path,verboseFalse,opset_versionopset,do_constant_foldingTrue,input_names[images],output_namesoutput_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_modelonnx.load(onnx_path)onnx.checker.check_model(onnx_model)print(✅ ONNX 模型结构验证通过)# ---- opset 版本 ----print(\n Opset 版本:)foropsetinonnx_model.opset_import:domainopset.domainifopset.domainelseai.onnxprint(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 上的效果还是集成到自动化模型转换流水线中这个方案都能显著降低成本、提升效率。欢迎直接拿去使用遇到问题可以在项目仓库中反馈讨论。