ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

RK3588部署RTMPose姿态估计:从ONNX到NPU量化的完整实战指南

RK3588部署RTMPose姿态估计:从ONNX到NPU量化的完整实战指南 最近一个项目要在RK3588板子上做实时人体姿态估计对比了一圈模型最终选了RTMPose。原因很简单RTMPose在COCO上精度够高推理速度在ARM CPU上也能跑到实时关键是它输出格式简洁配合RK3588的NPU做量化部署非常顺手。但真正踩进去才发现从MMPose导出ONNX到RKNN转换再到板端后处理每一步都有不少“文档里不会写”的细节。这篇文章就把我整个部署流程和踩坑记录完整捋一遍给后面做RK3588模型落地的同学一个能直接参考的实操路径。RTMPose出自MMPose仓库核心是基于SimCCSimultaneous Classification and Regression的坐标表示把关键点坐标预测转换成分类加回归的任务。相比传统的Heatmap方法RTMPose输出特征图小解码逻辑简单在低算力设备上非常友好。而RK3588的NPU理论算力6 TOPS支持INT8/INT16量化正好适合把RTMPose这种“轻量TransformerCNN”结构压到板上跑。这篇内容适合有基本深度学习基础、想在瑞芯微平台上做视觉应用尤其是姿态估计、动作识别方向的同学。1. RK3588的NPU与RTMPose的适配性为什么这个组合值得部署1.1 RK3588算力与内存布局对姿态估计的影响RK3588这块芯片最强的不是CPU而是内置的6 TOPS NPU。它的NPU支持INT4、INT8、INT16量化也可以跑FP16混合精度但实际工程中大家基本都用INT8因为算力利用率最高。内存带宽方面RK3588支持LPDDR4X/LPDDR5双通道带宽足够喂饱NPU实测下来不会出现“数据搬运等半天”的局面。但这里有个关键点RK3588的NPU对模型结构比较挑剔。它内部是三级流水线每个算子都有对应的硬件加速单元如果模型里出现某些特殊算子比如动态shape、某些高版本Transformer的LayerNorm变体就可能变成CPU回退速度直接崩掉。RTMPose整体结构是主干网络如CSPNeXt或ResNet SimCC头主干是卷积为主头是卷积PSS整体算子类型比较规整NPU支持度高这是它适合RK3588的根本原因。1.2 RTMPose各版本在边缘设备上的取舍RTMPose有s、m、l、x等多个版本还有针对不同主干网络的变体。在RK3588上部署我建议优先考虑RTMPose-s或RTMPose-m。s版本参数约5.6M在RK3588的NPU上量化后单帧推理大概能跑到20ms左右m版本精度更高但耗时大约翻倍。具体取舍要看项目需求——如果只是做人形检测和关键点可视化s就够用如果要做动作识别或康复评分建议直接m精度余量更足。另外要注意MMPose官方给的RTMPose权重是基于COCO训练的其中coco数据集80类里只用了人这一类关键点17点。导出模型的时候输出头的维度要和你的业务对齐比如你只要上半身6个点那需要重新训练或者改输出头不能直接拿原版权重的17点输出去截断否则后处理会对不上。2. 环境搭建RKNN-Toolkit2与MMPose的联动坑2.1 在PC端准备RKNN-Toolkit2环境RKNN-Toolkit2是瑞芯微官方提供的模型转换和仿真工具。先说明一下真正部署到板子上有两条常见路线路线A在PC上装rknn-toolkit2pip安装用Python API把ONNX转成RKNN然后在PC上用模拟器模拟NPU执行做精度验证再把RKNN文件丢到板子上通过板端rknn_toolkit_lite或librknnrt.so加载。路线B在板子上直接装rknn-toolkit2但板端内存和CPU资源有限转换大模型时容易爆内存而且校准数据集也要跟着板子走很麻烦。我只推荐路线A。PC端安装rknn-toolkit2时有个大坑它依赖的numpy、opencv版本必须和官方文档完全一致否则import就报错。我用的Python 3.8环境执行pip install rknn-toolkit2-1.5.0-cp38-cp38-linux_x86_64.whl装完后立刻验证from rknn.api import RKNN print(rknn ok)如果报ImportError大概率是numpy版本太新比如numpy 1.24需要降级到1.23.5。另外推荐在虚拟环境里装别直接用系统Python否则后面换项目很容易被依赖地狱坑到。2.2 用MMPose导出RTMPose的ONNX标准姿势之外的注意点MMPose从v1.0开始建议用mmdeploy导出ONNX但mmdeploy对RKNN的支持并不直接最稳定的方法是绕过mmdeploy直接基于PyTorch模型结构导出ONNX。原因是RTMPose的前向计算逻辑本身很简单不需要太多自定义算子。我用的是mmpose1.3.0加载训练好的RTMPose-m权重写一段自定义导出脚本。核心逻辑是import torch from mmpose.apis import init_pose_model model init_pose_model(rtmpose_m_8xb32-256x192.py, rtmpose_m.pth, devicecpu) model.eval() dummy_input torch.randn(1, 3, 192, 256) torch.onnx.export(model, dummy_input, rtmpose_m.onnx, input_names[input], output_names[simcc_x, simcc_y], opset_version11, do_constant_foldingTrue)这里必须提到两个问题opset_version不要太高RKNN-Toolkit2对opset 13以上的ONNX支持偶尔会有算子不兼容的问题另外RTMPose的输出实际上是两个SimCC分支分别是simcc_x和simcc_yshape为[1, K, num_bins]其中num_bins取决于输出分辨率比如192x192输出时num_bins就是192。导出后务必打印ONNX节点的输出shape确认和原模型一致我见到不少人导出后输出shape对不上后处理直接崩。2.3 板端运行库的版本匹配问题板子上运行RKNN模型时需要安装rknn-toolkit-lite2或者直接用librknnrt.so。这里有几个版本坑PC端转换时用的rknn-toolkit2版本必须和板端rknpu2驱动版本严格对应。比如我用1.5.0板端就要用librknnrt.so对应1.5.0的runtime。如果版本不一致加载模型会直接报错或者更诡异的是模型能加载但推理结果全是nan。板端系统如果是Debian/Ubuntu建议直接用官方编译好的rknpu2运行库别自己从源码编译耗时且容易缺头文件。如果板端用了Docker容器必须把/dev/rknpu设备映射进容器否则无法加载NPU。检查设备节点最简单的方式ls /dev/rknpu*如果不存在说明驱动没加载先检查dmesg是否有rknpu相关异常信息。3. 模型转换实战从ONNX到RKNN的完整流程3.1 用rknn.config配置量化与优化拿到RTMPose的ONNX后下一步就是转成RKNN。初始化RKNN对象并做配置代码框架固定from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[0.485, 0.456, 0.406]], std_values[[0.229, 0.224, 0.225]], target_platformrk3588, quantized_dtypew8a8, quantized_algorithmnormal, optimization_level3 )这里mean_values和std_values必须和模型训练时的预处理一致。RTMPose在MMPose中的预处理是归一化到ImageNet的均值和标准差所以这里不能写成常见的0-255归一化。很多新人踩坑就在这直接用了0-255的归一化结果模型转换成功但精度完全是乱的。quantized_dtypew8a8表示权重和激活都量化成INT8这是RK3588最常用的配置。如果对精度要求高可以试试w8a16但速度会打折。3.2 数据集准备与量化校准量化模型时需要一个校准数据集它是用来统计每层激活的数值范围所以数据要尽量贴近真实场景。我这边直接用COCO验证集里抽了200张人形图片预处理成和训练时一样缩放至256x192归一化。校准数据集的加载方式有两种通过rknn.load_onnx后调用rknn.build时直接传入dataset参数指向一个txt文件每行放一张图片的路径。或者在rknn.export_rknn后的推理验证阶段用rknn.init_runtime加载。实际操作我建议用第一种把校准数据和目标平台绑定在一起减少后续手动加载出错的概率。rknn.load_onnx(modelrtmpose_m.onnx) rknn.build(do_quantizationTrue, datasetcalib_data.txt)校准图片不要太多200张足以多了反而会让量化时间翻倍。校准图片太少比如10张则可能会让激活范围统计不准确导致量化后精度骤降。3.3 模型精度验证用py推理接口先跑通转换完成后建议先在PC端模拟器上验证精度再上板。用rknn.init_runtime(host_targetx86)在PC上模拟NPU执行import numpy as np from rknn.api import RKNN rknn.load_rknn(rtmpose_m.rknn) rknn.init_runtime(host_targetx86) # 读取一张真实图片做前处理 img cv2.imread(test_person.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img cv2.resize(img, (192, 256)) img img.astype(np.float32) / 255.0 img (img - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225]) img np.transpose(img, (2, 0, 1))[None, ...] outputs rknn.inference(inputs[img])这个outputs就是量化后在模拟器上的输出和PyTorch导出的ONNX输出对比一下差异。如果关键点坐标在5个像素以内基本说明量化质量可以接受。这里我要分享一个经验RKNN模拟器的结果和真实NPU结果可能存在微小差异尤其在某些LayerNorm实现上所以模拟器只是过滤掉大问题最终一定要以板端实跑结果为准。4. 板端推理代码实现从Python到C的迁移4.1 Python版推理读图、前处理、后处理一肩挑在板子上最简单的验证方式是Python调rknn-toolkit-lite2。先把rknn复制到板子然后from rknnlite.api import RKNNLite rknn_lite RKNNLite() rknn_lite.load_rknn(rtmpose_m.rknn) rknn_lite.init_runtime(core_maskRKNNLite.NPU_CORE_AUTO) img load_and_preprocess(test_person.jpg) outputs rknn_lite.inference(inputs[img])这里的core_mask可以指定用NPU的哪个核心RK3588有三个NPU核心默认NPU_CORE_AUTO会自动调度。如果你只有一个模型在跑设成NPU_CORE_0反而更稳定可以减少核心切换开销。Python版的主要作用是快速功能验证。它能跑通后我们就可以进入后处理和性能优化阶段。4.2 后处理细节RTMPose的解码与关键点可视化RTMPose的输出是SimCC的两个分支后处理本质上是把simcc_x和simcc_y在对应轴上做softmax然后求期望或者是取最大值得到关键点坐标。伪代码如下def decode(simcc_x, simcc_y): N, K, W simcc_x.shape # 对每个关键点在x和y方向做softmax probs_x softmax(simcc_x, dim-1) probs_y softmax(simcc_y, dim-1) # 计算期望位置 coord_x torch.sum(probs_x * torch.arange(W), dim-1) coord_y torch.sum(probs_y * torch.arange(W), dim-1) # 这里用期望比argmax平滑测试下来稳定性更好 return coord_x, coord_y这里有个细节RTMPose训练时用的是SimCC表示坐标范围是图像尺寸192或256所以解码后的坐标是相对于输入图像的要放缩回原始图像需要乘以比例系数(orig_w / 192, orig_h / 256)。还有SimCC的softmax输出可以当置信度用用于过滤低质量关键点。可视化结果时不要直接画在解码坐标上建议先用Affine变换把坐标映射回原图否则坐标对不上。最简单的办法是用np.linalg.inv做逆仿射变换。具体来说如果训练前处理是用仿射变换把图片缩放加裁剪到192x256那后处理就要用对应的逆矩阵。4.3 C部署时的内存与耗时优化功能验证通过后如果项目是落地到产品必然要写C部署。主要原因是Python解释器在板端占用CPU而且RKNNLite.inference会有一层Python调用的开销。C版核心流程是#include rknn_api.h rknn_context ctx; rknn_init(ctx, rknn_file, 0, 0, nullptr); rknn_input inputs[1]; inputs[0].type RKNN_TENSOR_FLOAT32; inputs[0].fmt RKNN_TENSOR_NCHW; inputs[0].buf input_data; inputs[0].size 3 * 192 * 256 * sizeof(float); rknn_output outputs[2]; outputs[0].want_float 1; outputs[1].want_float 1; rknn_run(ctx, nullptr); rknn_outputs_get(ctx, 2, outputs, nullptr);常见的内存优化手段是复用输入输出buffer避免每次推理都重新申请。RK3588的NPU对连续内存友好如果用rknn_query获取输入输出属性后用rknn_create_mem创建共享内存IO开销能降30%以上。实测用rknn_create_memrknn_set_io_mem比直接传CPU buffer快不少尤其在连续丢视频帧做实时推理时更明显。C里另一大块是耗时统计。用chrono包住rknn_run只统计NPU推理时间不包括前后处理。正常来说RTMPose-m量化后NPU推理时间大概25ms整个流程读图预处理推理后处理可视化在CPU单线程下大约60ms优化后可以到35ms已经可以做到实时。5. 实测效果与调优记录帧率、延迟、内存占用5.1 不同版本与量化配置的实测对比我在一块自己焊的RK3588板子上做了几组对比测试系统是Ubuntu 22.04Armbian内核默认CPU频率调到了performance模式。测试图片为一张1280x720的全身人物照输入尺寸256x192结果如下模型量化方式NPU推理耗时(ms)CPU后处理耗时(ms)CPU占用(%)RTMPose-sINT812.43.221RTMPose-mINT824.84.632RTMPose-mw8a1635.94.634RTMPose-lINT841.25.143这里w8a16表示权重INT8、激活INT16精度略高但耗时明显上升对实时项目不划算。RTMPose-l在RK3588上也能跑但帧率只能到24Hz左右如果做实时视频流就略吃紧。另外我还测了batch1和batch4两种情况batch4时NPU吞吐提升明显适合一次处理多路视频流的场景。如果你用rknn_run多次提交记得在rknn_set_inputs时把批量维设对。5.2 常见跑偏问题姿态抖动、关键点漂移部署中最容易遇到的两个古怪现象第一个是姿态抖动。明明模型转换验证时精度很高跑视频流时关键点却一跳一跳的。排查后发现问题在后处理——我把SimCC期望位置直接除以比例系数放缩回原图但没有做中心点对齐。RTMPose训练时输入图像会被仿射变换居中处理解码时需要还原这层改动否则坐标有固定偏移视频中看起来就是所有点都偏左下。修法就是严格记录训练时的仿射变换矩阵在解码时做逆变换。第二个是关键点漂移尤其手肘、膝盖这些点。这往往是量化精度不够导致的。临时解决办法是关掉optimization_level从3降到1精度会好一些但速度会慢。更彻底的办法是增加校准数据里的人物多样性特别是不同光照、远近距离量化时激活值范围更合理漂移会少很多。5.3 如何系统性地调优RK3588上的整体性能除了模型本身板端性能还受其他因素影响。我的调优顺序是CPU定频。默认系统可能是powersave模式NPU跑得再快前处理和后处理卡CPU也无济于事。先sudo cpufreq-set -g performance把CPU高频锁定实测整体延迟能降20%。NPU核心绑定。应用只有单路推理时可以用core_maskRK3588_NPU_CORE_0绑核心减少核心间调度抖动。多路并行时才用AUTO。CPU线程池。前处理是纯CPU操作把图像缩放改成cv::resize的多线程版本或者用OpenCV的UMat异步执行能让前处理耗时从8ms压到3ms。内存池复用。不要每帧都重新申请cv::Mat和rknn_output在流式计算里反复申请释放内存内存碎片会导致卡顿严重时掉帧。根据我实际调整后的流水线RTMPose-m在RK3588上CPU performance NPU CORE_0 内存复用整个处理链路从摄像头采集到关键点输出平均单帧耗时31ms能跑满32Hz左右完全能满足常规的轻量级动作分析和交互应用。最后再分享一个我自己常备的小技巧在板端跑模型时用perf top或者top -d 1盯着看如果发现python或加载进程CPU占用超过50%大概率是后处理或数据拷贝在做无用功优先检查是不是用了Python列表循环而没有走numpy向量化。这类问题优化完整个体验会有质的提升。
返回列表