十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

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

RK3588上部署RTMPose人体姿态估计实战:从ONNX到RKNN全流程

RK3588上部署RTMPose人体姿态估计实战:从ONNX到RKNN全流程 前阵子接了个边缘计算盒子的项目要在RK3588开发板上跑人体姿态估计模型选了OpenMMLab家的RTMPose。这活儿乍一看不算难——模型现成、板子现成、RKNN工具链也都公开但真要把RTMPose部署到RK3588的NPU上跑起来中间还是有不少东西需要理顺。这篇文章就把我完整的部署过程、踩过的坑、以及最终验证过的方案分享出来给准备在RK3588上跑姿态估计的朋友做个参考。先说清楚这篇文章能帮你解决什么如果你手上有一块RK3588开发板想在本地用NPU跑RTMPose做单人或多个人的关键点检测不希望自己在算子报错、量化精度下降、前后处理不匹配这些环节里反复折腾那这篇内容会比较适合你。我会从模型选型聊到ONNX导出、RKNN转换、板端推理再到性能调优和问题排查尽量把每一步“为什么这么做”也讲明白。1. 部署前的硬件认知与方案选型1.1 RK3588的NPU到底该怎么用RK3588是瑞芯微的旗舰级SoCCPU部分用的是4×Cortex-A76加4×Cortex-A55性能在同档位板子里算是很能打的。但做模型部署真正要关注的是它的NPU官方标称算力是6 TOPS支持INT4、INT8、INT16这些量化精度配合RKNN工具链使用。这6 TOPS看起来不大但在边缘设备里跑RTMPose这种轻量级姿态模型算力是够用的。这里有个容易踩的认知坑很多人以为把模型丢到NPU上就自动加速了其实RK3588的NPU只认RKNN格式的模型不能直接加载ONNX或者PyTorch的权重。你需要先在x86的宿主机上用RKNN-Toolkit2把模型转换好再把转换出来的.rknn文件放到板子上加载推理。这套流程和你之前接触过的TensorRT部署思路很像但工具链、算子约束、调试方式都不太一样所以最好先花点时间把RKNN的转换和板端推理接口搞清楚。我用的板子是带8GB内存的RK3588开发板系统刷的是Ubuntu 20.04桌面版aarch64这套流程换成其他RK3588方案也基本通用。板端推理有两种方式一种是直接用Python版本的rknn-toolkit-lite2适合快速验证效果另一种是C/C接口librknnrt.so适合做正式的产品集成。我的建议是先用Python把整个推理流程跑通再决定要不要搬C毕竟前面模型转换和精度验证才是真正费时间的部分。1.2 RTMPose模型家族怎么选RTMPose是OpenMMLab旗下MMPose项目里的一套实时姿态估计模型主打“高精度低延迟”结构上用的是CSPNeXt骨架加SimCC式关键点头。和传统的Heatmap-based方法相比RTMPose把关键点坐标建模成坐标分类问题输出的是沿x、y两个方向的分类概率后处理更简单在端侧部署时也更友好。RTMPose家族按模型大小分成RTMPose-t、RTMPose-s、RTMPose-m、RTMPose-l这几个档次。我这次选的是RTMPose-m输入尺寸是256×192在COCO数据集上的精度大约在75 AP左右具体数值跟训练配置有关。如果你的板子对延迟更敏感可以换RTMPose-s速度更快精度会略降一点如果追求最高精度可以上RTMPose-l但NPU上的耗时也会相应增加。还有一个需要考虑的点RTMPose本身是针对单人的姿态估计模型也就是说如果你要处理画面里有多个人的情况需要先用目标检测模型把人框出来再把每个检测框送到RTMPose里做关键点检测。这种“检测姿态”的top-down方案是当前最常用的做法但部署的时候要同时管理检测模型和姿态模型两个RKNN模型。如果你只需要单人人脸或单人体姿态其实可以跳过检测环节直接把裁剪好的单人图片送进去。2. 模型导出从PyTorch到ONNX2.1 用MMDeploy导出ONNX模型RTMPose模型一般是在MMPose框架下训练的权重的格式是.pthRKNN-Toolkit不能直接吃PyTorch权重所以我们要先把模型导出成ONNX格式。最省事的方式是用MMPose自带的部署工具或者MMDeploy来导出而不是自己手动搭模型结构再转权重那样很容易在算子层面出问题。我的操作步骤大概是这样的先创建Python虚拟环境安装好mmpose和mmdeploy然后找到模型对应的配置文件和权重文件。例如我用的是rtmpose-m配置文件在mmpose的configs目录下权重文件从OpenMMLab的模型库里下载。接着执行部署导出命令:python tools/deploy/export_pose.py \ configs/body_2d_keypoint/rtmpose/coco/rtmpose-m_8xb256-420e_coco-256x192.py \ rtmpose-m_simcc-coco_pt-aic-coco_256x192.pth \ --input-shape 1 3 256 192 \ --work-dir work_dir/rtmpose_m_onnx这里有个值得注意的点--input-shape参数我建议直接固定成1×3×256×192不要用动态维度。因为RKNN转换时如果输入维度是动态的NPU内部在某些算子上会走通用逻辑性能会打折扣甚至有些算子转换起来更麻烦。固定的shape对边缘部署来说更稳定也比动态shape更容易做性能优化。导出完成后work_dir里会生成一个end2end.onnx文件。我习惯先看一眼这个模型的输入输出信息用Netron打开或者用Python代码打印确认输出节点是不是两个SimCC分支。RTMPose的输出不是单个heatmap而是一个包含坐标分类的列表常见情况是输出两个tensor分别对应x方向和y方向的分类结果。你如果看到的是这种结构那说明导出没问题如果输出只有一个节点多半是导出配置里没有把两个branch都带出来。2.2 导出后的预处理与输入输出确认模型导出只是第一步导出之后必须把输入侧的预处理逻辑对齐。RTMPose在训练时的预处理是图像先resize到256×192然后做归一化归一化的mean和std用的通常是ImageNet的统计值mean[0.485, 0.456, 0.406]std[0.229, 0.224, 0.225]通道顺序是RGB。这个预处理信息非常关键因为RKNN转换时可以直接把归一化系数写进模型里让NPU在处理输入数据时自动完成归一化这样板端的CPU就不需要额外做一遍浮点除法能省掉不少耗时。后面讲RKNN转换的时候我再详细说怎么设置这几个参数这里先记住一件事训练和部署的预处理必须完全一致否则模型精度会明显下降而且这种下降非常难排查因为你看起来“什么都对”但结果就是不准。输出侧也要确认清楚。RTMPose-m在256×192输入下SimCC分支输出的维度大概是1×(17×2)×192×2之类的东西COCO数据集是17个关键点每个关键点有x和y两个方向不同版本的MMPose导出形式可能略有差异。我后来的做法是在导出后用ONNX Runtime在PC上先跑一遍单张测试图把输出结果和PyTorch原始模型的输出做一个对比简单核对一下量级和形状确认没问题再去转RKNN。这一步能帮你把“ONNX导出问题”和“RKNN转换问题”分开定位省下很多排查时间。3. RKNN模型转换与量化3.1 RKNN-Toolkit2环境搭建RKNN的模型转换是在x86宿主机上完成的用的是RKNN-Toolkit2这个Python包。安装的过程官网文档写得比较清楚核心要求是Python版本和依赖库要匹配。我当时是在一台Ubuntu 20.04的x86电脑上配的Python用了3.8装了rknn-toolkit2的wheel包然后顺手把torch、onnx、onnxruntime这些也装好方便后面做对比验证。环境装好后可以先跑一下import rknn验证是否安装成功如果报依赖缺失就按照提示补装对应版本的numpy、protobuf这些库。这里有个经常遇到的问题RKNN-Toolkit2对numpy版本有一定要求新版本的numpy有些接口变化会导致rknn报错如果你在安装后import失败先检查numpy是不是被装成2.x了。板子端装的是rknn-toolkit-lite2它是一个更轻量的包只负责推理不负责模型转换。宿主机是x86架构板子是aarch64架构两个环境里的wheel包不能通用装的时候一定要看好平台。我个人习惯在宿主机和板子端都记录下来装的具体版本号方便以后复现环境。3.2 ONNX转RKNN完整代码模型转换的完整代码并不长核心流程是加载ONNX模型配置量化数据集设置输入参数的归一化方式然后构建RKNN模型并导出。下面是我转换RTMPose时用的脚本代码里的注释我把关键参数都标出来了from rknn.api import RKNN rknn RKNN(verboseTrue) # 1. 配置量化与预处理参数 rknn.config( mean_values[[123.675, 116.28, 103.53]], std_values[[58.395, 57.12, 57.375]], target_platformrk3588, quantized_dtypew8a8 ) # 2. 加载ONNX模型 ret rknn.load_onnx(modelend2end.onnx) assert ret 0, load onnx failed # 3. 构建RKNN模型需要提供量化数据集 ret rknn.build( do_quantizationTrue, dataset./dataset.txt ) assert ret 0, build rknn failed # 4. 导出rknn模型 ret rknn.export_rknn(./rtmpose_m.rknn) assert ret 0, export rknn failed rknn.release()注意这里我在config里直接写了mean_values和std_values这样RGB数据在进入NPU之前就会被减去均值并除以标准差。mean和std的值我先按ImageNet统计值来的实际上要看你训练时用的预处理配置最好从mmpose的config里查一下保持一致才是最重要的。3.3 量化数据集与归一化设置量化是RKNN部署里最影响精度的一环。RKNN默认使用INT8量化它需要一组真实图片来统计每一层的激活值范围这组图片就叫量化数据集。dataset.txt文件里每一行写一张图片的路径我通常放100到200张有代表性的图片内容尽量覆盖你实际场景中会出现的情况——比如做人体姿态就多放一些不同姿态、不同光照、不同背景的人体图片。如果量化图片数量太少或者场景单一模型精度掉得会非常厉害经常会出现所有关键点全都连在一起或者置信度极低的情况。我在第一次转换最快速度踩过这个坑只放了20张测试图上去结果关键点坐标全飘了一度以为算子出问题。后来把量化集扩到150张不同场景的人体图精度才恢复正常。另外quantized_dtype我选的是w8a8也就是权重和激活都用INT8量化这是性能最好的配置如果精度实在拉不回来可以尝试w16a16或者用混合量化但推理速度会慢一些。使用RKNN-Toolkit2转换过程中还有一个比较常见的报错是“op not support”意思是ONNX里的某个算子RKNN工具链不支持。RTMPose这个模型大部分算子都能转但万一遇到不支持的算子优先检查模型里有没有比较冷门的自定义层。解决办法是尽量用MMPose自带的标准结构重新导出模型不要去修改模型的head部分。如果还是不行可以考虑把不支持的算子留在CPU上跑用RKNN的op融合或自定义算子功能来处理但那是进阶玩法新手不建议一上来就碰。4. 板端推理用Python快速验证4.1 安装rknn-toolkit-lite2模型转换好之后把.rtmpose_m.rknn文件拷贝到RK3588开发板上然后在板子上装rknn-toolkit-lite2。这个包在板端的安装也很直接通过pip安装对应aarch64的wheel包就行装好后写个简单的推理脚本验证加载和推理是否正常。板端的Python推理接口和宿主机上转模型时用的接口不是同一个。转模型时用的是RKNN类推理时用的是RKNNLite类两者的API相似但不完全一样注意别搞混。官方文档虽然都写了但我看到很多人包括我自己一开始都用错直接把宿主机上的推理代码搬到板子上跑结果第一行就报错。4.2 推理代码加载模型、预处理、SimCC解码板端推理脚本的逻辑是读取图像将图像resize到256×192转换颜色空间为RGB然后传入NPU推理。因为我在转换RKNN时已经设置了归一化参数所以板端代码里不需要再做减均值除方差的操作只需要把数据格式整理成模型期望的输入格式即可。下面是我完整跑通过的Python推理示例import cv2 import numpy as np from rknnlite.api import RKNNLite rknn_lite RKNNLite() ret rknn_lite.load_rknn(./rtmpose_m.rknn) assert ret 0, load rknn failed ret rknn_lite.init_runtime(core_maskRKNNLite.NPU_CORE_AUTO) assert ret 0, init runtime failed img cv2.imread(test.jpg) img cv2.resize(img, (192, 256)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 推理 outputs rknn_lite.inference(inputs[img]) # 根据RTMPose的SimCC输出解析关键点 # 这一步要看ONNX导出时的输出节点顺序我这里假设有两个输出x和y simcc_x outputs[0] # shape: [1, K, W] simcc_y outputs[1] # shape: [1, K, H]这里需要特别说明后处理部分。RTMPose不是直接输出坐标而是输出每个关键点在x和y方向上的分类得分你要对每个关键点在分类维度上找最大值位置然后换算成对应到原图的坐标。如果后处理逻辑写得不对关键点位置就会错位这也是很多人在把RTMPose算法迁移到RKNN时最容易懵的地方。我建议你在PC上用ONNX Runtime先验证一遍自己的解码脚本确保逻辑正确再搬到板子上跑这样能少走弯路。如果只是验证功能把人体关键点画到图像上用cv2画点连线就行。先简单看个结果确认关键点位置基本准确再谈优化和C移植。别在第一步就追求完美因为RKNN整个链路里变量太多先把“能跑”跑通了后面再逐步把每一环的细节抠到位。5. 生产部署C接口与性能优化5.1 C接口SDK基本结构Python验证没问题之后如果只是自己做实验或快速出demo其实已经可以收工了。但如果是放到正式的项目里比如做边缘计算盒子或者工业产品那还是得用C接口把推理封装成服务。RK3588的C推理是基于librknnrt.so的头文件在rknn_api.h里接口风格和很多加速器SDK类似初始化上下文、查询输入输出属性、设置输入数据、运行推理、获取结果。一个最简的C推理流程大致是这样的#include rknn_api.h // 读模型文件 FILE *fp fopen(rtmpose_m.rknn, rb); fseek(fp, 0, SEEK_END); int model_len ftell(fp); rewind(fp); void *model_data malloc(model_len); fread(model_data, 1, model_len, fp); fclose(fp); // 初始化 rknn_context ctx; rknn_init(ctx, model_data, model_len, 0, nullptr); // 查询输入输出 rknn_input_output_num io_num; rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, io_num, sizeof(io_num)); // 创建输入tensor填充数据 rknn_input inputs[1]; inputs[0].index 0; inputs[0].type RKNN_TENSOR_UINT8; inputs[0].size 256 * 192 * 3; inputs[0].fmt RKNN_TENSOR_NHWC; inputs[0].buf img_data; rknn_inputs_set(ctx, 1, inputs); // 运行推理 rknn_run(ctx, nullptr); // 获取输出 rknn_output outputs[2]; outputs[0].want_float 1; outputs[1].want_float 1; rknn_outputs_get(ctx, 2, outputs, nullptr);C这一层没有太多神乎其神的技巧核心就是严格按照SDK的接口把数据填对。最容易出问题的地方是输入数据的排布RK3588的NPU对NHWC格式支持得更好所以如果你的图像是HWC的内存布局直接填进去就行不要画蛇添足转成CHW。拿到输出之后SimCC解码的逻辑和Python版本完全一样写好之后把解码结果通过回调函数传给业务层即可。为了省时间我一般会把图像读取、resize、推理、解码封装成同一个函数业务层只需要传入图像数据拿到的就是关键点坐标数组。5.2 性能摸底与优化方向模型能跑通之后最大的问题就是性能。RK3588的NPU虽然算力不错但实际耗时和模型大小、输入尺寸、量化方式、甚至NPU核心分配方式都有关系。我第一次用RTMPose-m加上256×192输入做全INT8推理测出来单帧耗时在80ms左右也就是12FPS左右对于实时姿态估计场景来说还不够顺滑。优化点主要有这么几个方向按照投入产出比我来排个序第一优先确认是否启用了NPU加速而不是CPU兜底。RKNN推理如果某些算子没转成功SDK会自动把算子落到CPU上跑推理耗时就会急剧上升。你可以在init_runtime时设置打印日志看看有没有算子落到CPU上执行。第二合理设置NPU核心。RK3588的NPU有三个核心RKNNLite初始化时可以用NPU_CORE_0/1/2或NPU_CORE_AUTO指定使用核心我测下来NPU_CORE_AUTO在默认情况下不一定最优手动绑定到某个核心有时候还更稳定。不过让多路模型并发跑在不同核心上是更好的利用方式这个要看你的具体业务。第三如果精度允许可以换更小的模型或更小的输入尺寸。RTMPose-s配192×192的输入通常比RTMPose-m配256×192快30%到50%对于单人或简单姿态场景精度差距并没有想象中那么大。如果你不需要特别精细的关键点直接把输入分辨率降一档效果非常明显。第四考虑前后处理优化。resize和颜色转换如果放在CPU上做也会吃掉不少耗时。对于视频流场景可以尝试用板子的RGA硬件做缩放和格式转换把resize从CPU解放出来。这个优化比较高级但效果显著实测下来单帧整体耗时能再降十毫秒以上。经过这些优化我最终在自己项目里把单帧姿态估计的耗时稳定在了40ms左右也就是25FPS的水平对大多数姿态识别场景来说已经足够用了。5.3 多路视频与多模型并发时的额外建议如果你的场景是要同时处理多路人体的检测加姿态估计那就要涉及“检测模型RTMPose模型”的协同调度了。RK3588的NPU调度有一个特点单次推理往往不能跑满三个核心所以多路模型并发时反而更容易把NPU的整体算力用起来。比如一个视频流里检测模型和姿态模型交替推理天然就会把NPU的占用率拉高。我的实践做法是把检测模型和姿态模型分别初始化成独立的RKNN上下文然后放进一个推理线程池里统一调度。线程池一方面控制推理的并发数另一方面避免频繁创建销毁上下文带来的开销。这里有个细节初始化RKNN上下文比较耗时而且会占用不少内存所以不要在每个视频帧里重复初始化一定要在程序启动时初始化好运行时只做推理。6. 常见问题与排查技巧6.1 问题汇总速查表我把这段时间遇到以及身边朋友问过比较多的RK3588部署RTMPose问题整理成了下表方便大家对照排查现象常见原因解决思路模型加载失败提示init runtime fail板端driver版本与宿主机RKNN-Toolkit不匹配检查板端librknnrt版本升级或降级rknn-toolkit-lite2推理结果全为0或NaN输入数据格式错误或归一化重复设置确认输入是NHWC还是NCHW确认是否又做了一遍归一化关键点位置偏移但PC端ONNX推理正常板端预处理与训练不一致检查resize方式、通道顺序、归一化参数精度大幅下降量化数据集太少或场景单一补充100~200张覆盖实际场景的量化图片重新build推理速度特别慢部分算子落到CPU执行打开RKNN日志查看哪些op用了CPU尝试算子替换或混合量化C编译报头文件不存在没找到rknn_api.h或librknnrt.so安装完整SDK编译时加-I和-L路径动态shape导致转换失败导出ONNX时没有固定shape重新导出ONNX输入尺寸设为1×3×256×192这种固定值多个模型切换后内存持续增长推理上下文没有释放或重复初始化将模型初始化放到启动阶段只初始化一次用上下文复用代替反复创建6.2 自己踩过的几个坑第一个坑是量化数据集和验证集混用。我之前图省事直接用验证集里的图片做量化结果看起来精度特别高但一到真实场景就翻车。因为量化时模型已经“看过”这些图片了评估结果会偏乐观。正确的做法是单独准备一组和实际应用场景相近的图片做量化量化完后再用另一组没有见过的图来验证精度。第二个坑是ARM板端Ubuntu系统自带的Python环境比较老直接pip安装rknn-toolkit-lite2时经常出现依赖冲突。我的建议是给板子创建一个干净的Python虚拟环境单独安装推理依赖不要污染系统环境。这样即使环境坏了删掉重建也就一分钟的事。第三个坑是C工程里的模型路径问题。很多人习惯用相对路径加载rknn模型文件但服务程序挂在系统服务下运行时当前工作目录往往不是项目目录模型加载就会突然失败。后来我统一改成绝对路径并在初始化时检查模型文件是否存在这种问题就彻底消失了。第四个坑是关于摄像头输入。如果你直接用USB摄像头抓帧OpenCV默认输出BGR而RTMPose模型接收的是RGB需要转换。漏掉这一步的话关键点检测的准确率会掉得非常厉害尤其是颜色敏感的动作比如穿红衣服时手部检测会明显漂移。7. 一些后续可以继续做的事情部署做完之后如果你想让这个姿态估计方案的稳定性和扩展性更好几个方向你可以继续研究一下。一是把姿态估计和业务逻辑解耦。目前我们默认的是“得到关键点坐标就完事”但实际项目里往往还有动作识别、姿态纠正、行为分析这些上层应用。把姿态估计封装成独立推理服务输出统一格式的关键点数据后面接什么业务都方便不用每次改算法都动整个工程。二是考虑多路并行和负载均衡。如果你有多个视频源要同时处理可以在一台RK3588上同时跑多个进程或线程每个进程绑定不同的NPU核心这样资源利用更充分。需要做好内存规划和线程同步避免推理时互相干扰。三是模型继续迭代。RTMPose本身是在公开数据集上训练的如果实际场景里的人员姿态和训练集差异比较大比如全是侧身、低头或者穿宽大服装的人建议在自有数据上做微调。微调之后重新导出ONNX、重新转RKNN就能用。我在实际项目中的最大体会是RK3588部署RTMPose这件事本身并不算难真正麻烦的是整个链路里那些看似不起眼的细节——量化数据集是否贴合场景、预处理是否和训练一致、输入输出格式对不对、NPU核心有没有用好。这些小细节单独拎出来都不起眼但每一个都能让推理结果从“看起来能用”变成“完全不能用”。按我上面这套流程走一遍你至少可以把这些雷区都避开然后再在速度和精度之间找到适合自己的平衡点。
返回列表