ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

RK3588模型部署实战:从ONNX到RKNN的完整转换指南

RK3588模型部署实战:从ONNX到RKNN的完整转换指南 上个月接了个项目要把训练好的手势识别模型部署到瑞芯微RK3588板子上跑。模型在服务器上推理跑得飞快一上板子问题就来了rk3588的NPU不认识PyTorch的权重只认自己的rknn格式。折腾了好几天把整个转换链路跑通之后回头看其实核心就三件事把训练模型导成ONNX、用rknn-toolkit2转成rknn、在板端用rknnlite加载推理。每一步都有不少坑尤其是算子兼容性和量化精度这两块踩得我一度怀疑人生。这篇文章把这套流程完整记录下来从环境准备、ONNX导出、RKNN转换参数配置到板端部署代码全是我调试通过的真实方案。想做rk3588模型部署的朋友不管是手势识别、YOLO目标检测还是其他视觉模型这套流程基本都能直接用重点看第二章的ONNX导出规范和第三章的量化配置心得。1. 整体思路与方案选型1.1 为什么选择RKNN工具链而不是直接在板子跑PyTorch先把结论放在前面rk3588上想用上NPU的算力模型必须经过瑞芯微官方的RKNN工具链转换这是唯一路径。虽然板子上的ARM CPU可以硬跑PyTorch或者ONNX Runtime但效果完全不是一个量级。rk3588的NPU标称6 TOPS算力对应的是INT8精度下的卷积计算。CPU跑一个轻量级手势分类模型单帧推理可能需要50到100毫秒而NPU运行同样模型通常在5到15毫秒差距非常明显。更关键的是NPU跑模型时CPU占用率几乎为零可以腾出四核A76和四核A55去处理摄像头采集、画面渲染、通信协议这些任务。瑞芯微的模型转换工具叫RKNN-Toolkit2它对模型格式的支持情况是这样的PyTorch模型可以通过torch.onnx.export导出为ONNX再喂给RKNN-Toolkit2TensorFlow的SavedModel、TFLite也可以直接或间接转换Caffe模型同样兼容。实测下来最稳妥的中间格式就是ONNX原因后面细说。RKNN-Toolkit2的工作模式有模拟推理和板端联调两种。在PC上装好工具后可以直接加载ONNX并模拟NPU推理结果方便验证模型转换是否正确不需要板子就能排查大部分算子问题。板端部署时则使用精简版的RKNNLite接口运行时依赖只有librknnrt.so一个动态库干净利落。1.2 模型转换的整体流程框架一个完整的手势模型部署流程分为四个阶段模型准备、格式转换、板端适配、性能验证。模型准备阶段要确认模型结构、输入尺寸、预处理参数均值、方差、归一化系数准备好校准数据集。校准数据集是量化阶段用的别漏了。格式转换阶段在PC上装好rknn-toolkit2加载ONNX模型设置好target_platform为rk3588配置量化配置和数据集build成rknn文件。这个过程是整套部署的绝对核心。板端适配阶段把rknn文件拷贝到开发板安装rknn-toolkit-lite2运行时库用Python或者C接口写推理代码。Python接口开发快C接口适合追求极致性能的场景。性能验证阶段跑通推理之后要做的三件事对比NPU输出和CPU输出的结果一致性测量单帧推理耗时和内存占用用真实摄像头视频流做长时间稳定性验证。很多人前两步做完了就直接上线结果摄像头跑半小时就内存爆炸这就是第三步没做。提示实际项目里我强烈建议在动手转换之前先写一个脚本把训练时的预处理、模型输出解析画成流程图贴在最显眼的地方。RKNN转换时的mean/std配置必须和训练时完全一致这里错了后面的功夫全白费。2. 环境准备与模型导出2.1 PC端环境RKNN-Toolkit2安装实录RKNN-Toolkit2目前是Python包的形式发布支持x86_64 Linux环境。我在Ubuntu 20.04上安装Python版本用了3.8其他常见组合是Ubuntu 18.04加Python 3.6或者Ubuntu 22.04加Python 3.10。不同版本的rknn-toolkit2对Python版本有约束装之前一定要看官方文档的版本对照表。安装步骤相当简单用pip就行# 建议用虚拟环境避免污染系统Python python3 -m venv rknn_env source rknn_env/bin/activate pip install rknn-toolkit2-1.6.0-cp38-cp38-linux_x86_64.whl依赖包会自动装包括numpy、onnx、Pillow、opencv这些基础库。装完验证一下python -c from rknn.api import RKNN; print(RKNN.__version__)如果这个命令不报错说明工具链装好了。要注意的是1.x版本之间API差异不大但部分算子支持和量化策略有细微区别同一个rknn文件在不同版本工具下生成的性能可能不同。正式项目建议锁定一个版本不要频繁升级。2.2 训练模型导出ONNX的标准化操作我的手势识别模型是用PyTorch写的结构不算复杂一个轻量卷积神经网络输入是3通道的RGB图像尺寸是224x224输出是5个手势类别。导出ONNX的代码基本上是这个模板import torch model GestureNet() model.load_state_dict(torch.load(gesture.pth, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, gesture.onnx, input_names[input], output_names[output], opset_version11, dynamic_axesNone )这里有两个非常关键的细节。第一opset_version我用了11不要一上来就选最新的17、18虽然新版本算子支持更全但RKNN工具链对高版本opset的兼容性可能滞后反而容易报错。第二dynamic_axes我直接设为None也就是固定批量大小为1。手势识别推理是一次处理一帧图像不需要动态batch固定shape能减少转换时的麻烦。ONNX导出成功后可以用onnx库检查一下输入输出信息import onnx onnx_model onnx.load(gesture.onnx) graph onnx_model.graph print(input:, [i.name for i in graph.input]) print(output:, [o.name for o in graph.output])这一步输出的名字后面RKNN脚本里要用。我经常看到有人在这里踩坑PyTorch导出ONNX时没指定input_names导致输入节点名变成random tensor这种莫名其妙的名字RKNN加载时怎么都找不到对应的输入端口。注意关于BatchNorm层的处理。如果你的模型训练代码里有BatchNorm结构不用做任何特殊操作PyTorch导出ONNX时会自动把它融合到卷积层中ONNX中的BN实际上是以卷积Bias形式展开的。如果你看到ONNX模型里还有独立的BatchNorm算子那多半是导出的方式有问题RKNN转换时很可能会报不支持。2.3 校准数据集的前期准备量化是RKNN转换必须经历的一步除非你选择跑全FP16精度。INT8量化需要用一小批真实图片来计算每一层的动态范围这批图片就是校准数据集。校准数据集的选择直接影响量化后模型的精度我用的是300张从测试集里随机抽出的图像覆盖了不同光照条件、不同手势姿态、不同背景干扰。把这些图片整理到一个txt文件里每行一个文件路径这就是dataset.txt。实际上RKNN-Toolkit2支持两种量化配置方式这里先不展开第三章细说。实操心得校准图片不是越多越好。我有一次为了追求精度塞了2000张图进去结果build时间翻了十倍精度却没有明显提升。瑞芯微的文档建议50到200张就差不多够用了覆盖多样性的优先级远高于数量。3. RKNN模型转换实操3.1 转换脚本的完整解析环境准备好、ONNX导出成功之后就进入了整个项目最核心的环节用RKNN-Toolkit2把ONNX转换成rknn格式。下面是我调试通过的完整脚本我逐行讲解关键配置的含义。from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[127.5, 127.5, 127.5]], std_values[[127.5, 127.5, 127.5]], target_platformrk3588, quantized_dtypewgt_sym_i8, quantized_algorithmnormal, optimization_level3 ) rknn.load_onnx(modelgesture.onnx) rknn.build( do_quantizationTrue, datasetdataset.txt ) rknn.export_rknn(gesture.rknn) rknn.release()这个脚本看起来简单但几乎每一行都值得仔细琢磨。mean_values和std_values这一组参数作用是设置模型输入的归一化方式。我的模型训练时是这样预处理的把RGB像素从0到255归一化到-1到1也就是 (pixel - 127.5) / 127.5。那么配置里mean_values就是[[127.5, 127.5, 127.5]]std_values是[[127.5, 127.5, 127.5]]。这里如果填错了或者训练时归一化方式不一样模型跑出来的预测概率会乱得一塌糊涂但完全没有报错。有个小知识点如果你训练时没有做归一化只做Resize到224x224那mean_values填[[0,0,0]]std_values填[[1,1,1]]就行。如果你的归一化是ImageNet风格的 (pixel/255 - mean)/std那就填对应的值。3.2 量化配置精度与速度的平衡术量化配置里的quantized_dtype我选了wgt_sym_i8这是权重对称量化、激活值非对称量化的组合也是瑞芯微NPU上速度最快的模式。除了这个还有以下选项wgt_sym_i8INT8权重对称量化推理速度最快推荐首选wgt_hybrid_i8混合量化部分敏感的层保留FP16精度略高速度略慢wgt_asym_i8权重非对称量化对某些不服从零对称分布的权重有奇效fp16不量化全程FP16推理精度最高速度约为INT8的一半对于手势识别这种任务我的经验是直接用wgt_sym_i8如果精度不达标再考虑hybrid或者fp16。不需要过度设计INT8没你想的那么脆弱。quantized_algorithm我选了normal即常规量化算法。如果量化后精度下降比较多可以试一下mmse最小均方误差算法它对动态范围的估计更精细但build时间会显著变长。实测下来mmse对某些模型能提升0.5到2个百分点的准确率。optimization_level填3表示做最大程度优化。RKNN工具链会自动做层融合、算子替换、内存复用。手势模型这种规模几百KB级别的在线和离线两种模式差异不大如果是几MB以上的大模型离线模式build时间更短但可能需要更多系统内存。重要提示量化数据集dataset.txt的路径必须是绝对路径或者相对于当前工作目录的相对路径。我第一次写的时候用了Python脚本所在目录的相对路径结果rknn.build一直找不到文件也不报错就是输出空结果后来排查了半天才发现是这个问题。3.3 模拟推理验证上板前的最后体检build成功导出gesture.rknn之后先别急着拷到板子上。RKNN-Toolkit2支持在PC上模拟NPU推理这一步能提前发现模型转换后是否存在数值异常。from rknn.api import RKNN import numpy as np from PIL import Image rknn RKNN() rknn.load_rknn(gesture.rknn) rknn.init_runtime(targetNone) img Image.open(test_gesture.jpg).resize((224, 224)) inputs np.array(img).astype(np.uint8) # 注意RKNN的输入通道顺序是HWC不是CHW outputs rknn.inference(inputs[inputs]) print(outputs)init_runtime不加任何参数就是模拟推理模式。这个模式不需要板子直接把rknn模型跑在PC上。我强烈建议在这个步骤做一次完整的验证把ONNX模型加载进来用同一张图片分别在ONNX Runtime和RKNN模拟推理上跑对比输出差异。两个输出在softmax之后的结果差异在2到3个百分点以内属于正常现象如果差得离谱重点检查mean/std配置和量化数据集。这里还有一个容易忽略的点输入数据的格式。很多人在PyTorch里习惯了CHW格式喂给RKNN的却是HWC导致模型输出完全不对。RKNN的Python接口默认接收HWC格式图像不需要做通道反转但一定要确认Resize尺寸和模型输入要求一致。上面代码里我特意用PIL的resize处理了图像如果你用OpenCV注意它的默认通道顺序是BGR如果你的模型在训练时用的是RGB需要先转通道。3.4 性能优化NPU核数与内存策略rk3588的NPU有三个独立核心理论上可以并行处理三个推理任务或者一个任务分配到三个核心上跑。但实际中绝大多数情况下一个手势模型单独用一个核心就足够剩余核心留给其他视频流或任务。在rknn.config阶段可以通过NPU_CORE配置核心使用方式。RKNN-Toolkit2的config方法里有个参数叫npu_core可选值为0单核、1双核、2三核、3自动选择。比如rknn.config( mean_values[[127.5, 127.5, 127.5]], std_values[[127.5, 127.5, 127.5]], target_platformrk3588, npu_core0 )为什么不需要多核因为模型小、推理快时多核通信和任务分配的开销可能反而超过并行收益。只有当单张图推理耗时超过30毫秒时才值得考虑多核。手势模型这种轻量级网络npu_core设为0单核跑就够了多核的效果我用过一次实测没有显著提升反而在连续推理时偶尔出现帧错乱。4. 板端部署与常见问题排查4.1 RKNNLite运行环境搭建rknn文件生成后下一步就是把模型部署到rk3588开发板上。板端运行环境和PC端工具链是两套不同的库PC端叫rknn-toolkit2板端叫rknn-toolkit-lite2核心差异如下对比项rknn-toolkit2rknn-toolkit-lite2运行位置PC x86_64板端 ARM64核心组件RKNN类支持转换推理RKNNLite类仅推理依赖库rknn运行时onnxtorch等librknnrt.so内存占用数百MB极低板端部署的具体操作是这样的先把rknn-toolkit-lite2的whl包拷贝到板子上pip安装同时确保板子的librknnrt.so版本和PC端工具链版本匹配。这个版本匹配问题非常隐蔽通常的表现是PC端模拟推理一切正常上板之后加载模型报RuntimeError: RKNN init failed板子日志显示rknn internal error。解决办法很简单使用rknn-toolkit2版本对应的librknnrt.so而不是随便去下载一个。我用的版本是1.6.0对应的库文件在板端需要到1.6.0的固件或库包里找。最好的验证手段是板端执行如下命令python -c from rknnlite.api import RKNNLite; rknn_lite RKNNLite(); print(rknn_lite.__version__)如果输出正常再用排除法加载失败的rknn模型在PC模拟端是否能正常加载。能加载说明板端库有问题不能加载说明rknn模型本身用处不对。4.2 板端推理代码从单帧到连续视频流板端推理代码我用的是RKNNLite接口这个类和RKNN几乎一样但去掉了模型转换相关的方法只保留加载和推理。一个最基础的单帧推理逻辑from rknnlite.api import RKNNLite import cv2 import numpy as np rknn_lite RKNNLite() rknn_lite.load_rknn(gesture.rknn) rknn_lite.init_runtime(core_maskRKNNLite.NPU_CORE_0) img cv2.imread(test.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (224, 224)) img img.astype(np.uint8) outputs rknn_lite.inference(inputs[img]) scores outputs[0][0] pred_class int(np.argmax(scores)) print(predicted class:, pred_class)注意三个地方第一init_runtime时可以传入core_mask参数NPU_CORE_0表示用第一个核心。第二inference返回的是列表每个元素对应一个输出但输出是不是softmax后的概率取决于你的模型结构。如果模型只在训练时用了CrossEntropyLoss不一定会带Softmax层那推理时需要对输出做softmax操作。第三这里的img是原始uint8图像不需要手动归一化因为归一化的mean/std已经在转换时烧进rknn模型里了你直接把原始像素喂进去就行。连续视频流的推理稍微复杂一点因为要处理摄像头采集和NPU推理的流水线问题。我实测下来最简单的方案是启动两个线程一个线程从摄像头读帧并放进队列另一个线程从队列取帧做推理。这里有一个坑队列大小必须固定否则长时间运行会积压大量未处理的帧内存一路飙到爆。import threading import queue import cv2 frame_queue queue.Queue(maxsize2) def camera_thread(): cap cv2.VideoCapture(0) while True: ret, frame cap.read() if ret: if frame_queue.full(): try: frame_queue.get_nowait() # 丢弃旧帧保证实时性 except queue.Empty: pass frame_queue.put(frame) def inference_thread(): while True: frame frame_queue.get() rgb cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) rgb cv2.resize(rgb, (224, 224)) outputs rknn_lite.inference(inputs[rgb]) # 输出处理... frame_queue.task_done()这里舍弃旧帧、保留新帧的做法很关键可以保证处理速度跟不上摄像头帧率时也不会越积越多始终处理最新画面。4.3 高频问题速查与排查实录把这次部署过程中遇到的所有问题整理成一个速查表给各位参考问题现象可能原因解决方案build时报onnx op不支持的错ONNX算子超出RKNN工具链支持范围切opset版本到11或12检查模型里是否有特殊算子如Gather、ScatterND转换成功但板端推理结果异常mean/std配置错误或量化数据集偏差大核对训练预处理参数用无量化(fp16)跑一次对比精度init_runtime时报版本不匹配librknnrt.so与实际NPU驱动不一致同步升级板端驱动库或使用板厂配套的rknn-toolkit2版本推理结果类别总是固定那个输出解析错误模型输出为logits或没有softmax检查模型结构必要时在输出后加softmax运行时内存持续增长帧队列积压或推理结果引用未释放限制队列长度及时释放帧对象引用图像颜色偏色导致误判OpenCV读取的是BGR直接喂给了要求RGB的模型推理前用cv2.cvtColor或np.flip转换通道最长的一次排查经历是模型转换成功、PC模拟推理也正常但板端结果总是第0类。后来发现是我在模型里用了BatchNorm推理时没设置model.eval()导致BN的running_mean和running_var没生效导出的ONNX里面就是随机初始化的BN参数。这个概率非常低但极其隐蔽建议导出前用torch.no_grad()包裹并显式调用model.eval()并且在导出后跑一遍ONNX Runtime验证数值是否合理。4.4 性能实测数据最后给大家看一组我这次项目的实测数据传感器到手后先做个基线免得后面调优没方向。环境是rk3588开发板手势模型输入224x224INT8量化单核NPU指标数值单帧推理耗时INT86.8毫秒单帧推理耗时FP1611.2毫秒单帧预处理耗时resizecvtColor2.1毫秒内存占用持续运行30分钟156MBCPU占用率推理线程约8%摄像头帧率30FPS稳定这个数据表现我还是挺满意的。用video采集推理的方式实测端到端延迟从摄像头画面到输出手势类别大约在40毫秒左右完全满足实时手势交互的体验要求。板子剩余的计算能力还能同时跑一路视频编解码和多个传感器数据采集。实操心得如果觉得单帧推理时间还是太长先看预处理和输出的时间占比。有些时候瓶颈不在NPU而在cv2.resize和cvtColor这种CPU操作。可以尝试把Resize操作交给更底层的硬件加速库或者把输入尺寸固定得更小一些比如从224x224降到160x160推理时间差不多能砍掉一半手势识别这种任务精度损失通常可以忽略。写在最后的几点经验这个项目做下来最大的体会是模型转换这件事本身并不难难的是对每个参数背后的机制理解到位。很多人一看到RKNN工具链脚本就往上套mean/std填错、量化数据集不规范、输出解析不对每一个小坑都会让你浪费好几天时间。我后来把整个转换流程固化成了一个标准模板导出ONNX前先跑一遍ONNX Runtime推理做基准、RKNN转换时先用fp16不量化跑通链路、再用INT8量化并对比精度、板端部署前先跑30分钟稳定性测试。这套流程帮我后续部署其他模型包括YOLOv8这类目标检测模型时基本一次就能跑通。最后再分享一个小技巧RKNN-Toolkit2在build结束时会输出一个模型分析报告里面包含每层的耗时占比和内存占用这个报告很多人忽略。把它截图存下来后续做性能优化时对照着看比盲人摸象高效得多。比如我在报告里发现模型最后一层全连接耗时占比奇高果断换成了全局平均池化单帧推理时间直接又压缩了1.2毫秒。这种细节文档里找不到只能靠自己多折腾。
返回列表