ARTICLE DETAIL

资讯详情

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

TensorRT部署RTMPose人体姿态估计:从ONNX导出到C++推理实战

TensorRT部署RTMPose人体姿态估计:从ONNX导出到C++推理实战 简介面向计算机视觉与算法部署工程师这是一份使用TensorRT对RTMPose人体姿态估计算法进行推理优化与落地部署的完整实战项目。内容从RTMPose工作原理与模型结构讲起深入解析TensorRT的层融合、精度校准、内核自动调优和多流并行等优化机制再完整演示模型导出、转换、校准、序列化及部署流程并附带实验验证环节通过推理速度、准确率与资源消耗等指标对比量化展示优化效果适合希望提升模型在NVIDIA GPU上实时运行效率的研究者与工程师参考。压缩包共16个文件大小63.6MB以C源码cpp/h和Visual Studio工程文件为主包含engine推理模型、Python脚本及项目说明文档结构清晰便于对照阅读、编译与二次修改。已有240人学习下载可作为算法部署入门到进阶的实战案例。1. 用 TensorRT 跑 RTMPose这个算法部署项目为什么值得拆人体姿态估计在直播、体育分析和人机交互里越来越常见但真正决定它能不能落地的不是模型在 GPU 上跑得多准而是从 PyTorch 训练权重到 TensorRT 推理引擎这条路走不走得通。我见过太多团队卡在 ONNX 导出后掉点、TensorRT 加速不明显、C 推理代码跑不通这三个坎上。这个名为“算法部署-使用TensorRT部署RTMPose人体姿态估计算法”的项目恰好把这整条链路完整走了一遍从 RTMPose 模型导出 ONNX到 TensorRT 引擎构建再到一个可编译的 Visual Studio 工程rtmpose_tensorrt.sln里面既有 RTMDet 检测器做人体框提取也有 RTMPose 本身的关键点推理还附带了推理封装 inference.h/cpp 和 utils 工具。对于想把姿态估计算法部署到自有 C 服务里、又不想从零踩一遍 TensorRT 坑的工程师这份资源可以省掉大量试错时间。我拆完整个工程后发现它的价值不只在于“能跑”更在于它把后端封装、图像预处理、后处理解耦得比较清楚适合拿来改成自己的业务代码。2. RTMPose 与 ONNX 导出模型结构、导出命令与检查清单2.1 RTMPose 的模型结构和部署特点RTMPose 是上海人工智能实验室提出的实时姿态估计算法核心亮点是把 heatmap 回归方式改成了 SimCC。熟悉传统姿态估计的读者知道以前的做法是让模型输出一张高斯热力图每个关节对应一个概率分布然后从分布里取峰值坐标作为关键点位置。这种方式精度不错但存在一个天然问题输出层分辨率低比如从 256×192 输入降采样到 64×48 输出那么一个像素的误差在原始图上就放大了 4 倍。RTMPose 换了一种思路把坐标预测拆成两个方向的一维分类问题——水平方向预测 x 属于哪个区间垂直方向预测 y 属于哪个区间再通过期望值计算得到亚像素精度的坐标。这就是 SimCCSimultaneous Classification and Coordinate Regression的核心思想。由于输出不再需要 2D 特征图RTMPose 的网络尾部可以用更轻量的分类头替代整体推理开销比传统 heatmap 方法小得多。再加上 MAE 预训练加知识蒸馏的加持RTMPose 在 COCO 验证集上的 AP 和推理速度都相当能打。从部署角度看RTMPose 的模型结构完全由常规卷积、BatchNorm 和全连接层组成没有动态形状分支和自定义算子这给 ONNX 导出和 TensorRT 转换带来很大便利。对比一些用可变形卷积或者自定义池化的模型RTMPose 的导出几乎不会在算子映射阶段翻车这也是这个项目选它做算法部署示范的原因之一。2.2 从 PyTorch 权重导出 ONNX完整步骤与输出参数项目实战的第一步是把训练好的 RTMPose PyTorch 模型转成 ONNX。这一步看似简单但在导出参数上其实有几个值得扣的细节。常见的导出脚本逻辑如下import torch from rtmpose import RTMPose # 假设这是项目中的模型定义 model RTMPose(num_joints17, simcc_split_ratio2.0) checkpoint torch.load(rtmpose_tiny_distilled.onnx, map_locationcpu) model.load_state_dict(checkpoint[state_dict], strictTrue) model.eval() # 固定输入尺寸BGR 通道顺序RGB 和 BGR 的顺序会直接影响导出后的推理结果 dummy_input torch.randn(1, 3, 256, 192) torch.onnx.export( model, dummy_input, rtmpose.onnx, input_names[input], output_names[simcc_x, simcc_y], dynamic_axes{input: {0: batch}, simcc_x: {0: batch}, simcc_y: {0: batch}}, opset_version13, do_constant_foldingTrue, )这里我把simcc_split_ratio显式写了出来它控制 SimCC 输出头的宽度——比输入分辨率放大多少倍。项目里不同规模的模型这个值可能不同需要跟训练配置保持一致。opset_version选 13 是为了在 TensorRT 里有更好的算子兼容性一些老代码用 11我遇到过后续某些优化层无法识别的场景所以建议直接上 13 或 17。对于输出头RTMPose 导出的不是 N×17×H×W 的 heatmap而是 N×17×(M×K) 的两个预测向量。上述代码中simcc_x和simcc_y分别是水平方向和垂直方向的分类向量。这个理解很关键如果沿用 heatmap 的后处理逻辑去解码结果必然是错的。2.3 导出后的模型结构检查清单导出 ONNX 后不要急着转 TensorRT先做两件事用onnx.checker做静态校验再用onnxruntime跑一遍同一张输入检查数值是否正常。import onnx import onnxruntime as ort import numpy as np onnx_model onnx.load(rtmpose.onnx) onnx.checker.check_model(onnx_model) session ort.InferenceSession(rtmpose.onnx, providers[CPUExecutionProvider]) input_tensor np.random.randn(1, 3, 256, 192).astype(np.float32) outputs session.run(None, {input: input_tensor}) print(len(outputs), outputs[0].shape, outputs[1].shape)如果输出头的维度不对问题基本出在simcc_split_ratio与模型定义不一致或者导出的dynamic_axes设置不完整。另外建议检查一下输入的通道顺序项目里图像预处理按 BGR 读取并用std::vectorcv::Mat转 float 后归一化如果模型训练时用的是 ImageNet 的均值和方差推理端要保持一致不要想当然地采用 0~1 直接缩放。这个我在实际部署时吃过亏后面避坑章节会展开讲。3. TensorRT 的优化逻辑层融合、FP16 与序列化3.1 TensorRT 为什么能提速核心机制拆解TensorRT 的加速不是简单地把 ONNX 算子翻译成 CUDA 代码而是会针对 NVIDIA GPU 架构做全局优化。我拆这个项目的时候对它的层融合机制印象最深。比如常规的 ConvBatchNormReLU 结构在 PyTorch 里是三个独立算子在 TensorRT 中会被融合成一个 CBR 算子。这意味着原来需要三次 kernel 启动和三次显存读写现在只需要一次在 batch size 较小时这个收益尤其明显。另一个关键机制是多流并行。TensorRT 会将计算图切分成多个执行流让没有依赖关系的层在不同 CUDA stream 上并行执行。比如 RTMPose 中 SimCC 头部分成 x、y 两条支路理论上就是两个可以并行的子图。TensorRT 优化器会自动识别这种结构并分配流。对于 RTMPose 这类计算密集型模型我一般推荐至少打开 FP16 精度。FP16 的显存占用是 FP32 的一半计算吞吐在 Turing 及以后的架构上能提升数倍。项目中开启了 FP16 模式和 DLA如果显卡支持实际推理速度在 RTX 系列上提升显著。不过要注意FP16 可能带来轻微精度损失尤其是最后的 SimCC 分类输出如果对精度要求极高需要对比验证后再决定是否启用。3.2 用 trtexec 快速转引擎命令行参数与精度选择搭建 C 工程之前先用官方命令行工具trtexec快速把 ONNX 转成 TensorRT 引擎能提前发现大部分转换问题。命令如下/usr/src/tensorrt/bin/trtexec \ --onnxrtmpose.onnx \ --saveEnginertmpose_fp16.engine \ --fp16 \ --minShapesinput:1x3x256x192 \ --optShapesinput:1x3x256x192 \ --maxShapesinput:8x3x256x192 \ --workspace4096--fp16表示开启半精度--minShapes、--optShapes、--maxShapes定义了动态 batch 的边界其中optShapes是 TensorRT 做内核优化时的参考形状建议按实际生产中最大概率出现的 batch 设置--workspace指定了构建时允许使用的显存上限单位是 MB。如果显存吃紧可以适当降低这个值但构建时间会变长。转换过程中常见的报错有两类一是某个算子不支持——解决方案是升级 TensorRT 版本或回退 opset二是形状推导失败——通常是 model 导出时 dynamic_axes 没写全。如果trtexec能顺利跑通后面 C 工程的路就顺畅多了。3.3 引擎序列化与反序列化避免重复构建的耗时TensorRT 引擎构建是个耗时操作一张 RTMPose 的 ONNX 模型在桌面级 GPU 上构建可能花上几分钟在嵌入式设备上更久。因此把 engine 序列化到磁盘以后每次启动直接反序列化使用是部署时必须要做的事。// 保存引擎到文件 std::ofstream file(rtmpose_fp16.engine, std::ios::binary); file.write(engine-serialize()-data(), engine-serialize()-size()); // 反序列化引擎 std::ifstream file(rtmpose_fp16.engine, std::ios::binary); std::vectorchar data(std::istreambuf_iteratorchar(file), {}); auto runtime nvinfer1::createInferRuntime(logger); auto engine runtime-deserializeCudaEngine(data.data(), data.size());这里有个反直觉的坑TensorRT 引擎绑定了构建它的 TensorRT 版本、GPU 架构和 CUDA 版本。在一台机器上构建的 engine拷贝到另一台相同显卡但驱动不同的机器上很可能无法反序列化。所以实际部署时更稳妥的方案是部署机器上先装好同一版本的 TensorRT再执行引擎构建或反序列化流程。项目里提供了模型转换脚本目的就是让用户能在自己的环境上重新生成引擎。4. 搭一个可跑的 C 推理工程RTMPose TensorRT 的完整实现4.1 工程结构梳理sln 解决方案里各文件的职责打开rtmpose_tensorrt.sln整个工程的模块划分相当清晰。我把关键文件的职责整理如下文件职责inference.h / inference.cpp推理封装层负责 TensorRT 引擎加载、输入输出 buffer 管理rtmdet.h / rtmdet.cppRTMDet 检测器推理逻辑用于检测人体框rtmpose.hRTMPose 关键点推理的头文件定义关键点数量和 SimCC 解码接口main.cpp程序入口串联 RTMDet 检测和 RTMPose 姿态估计utils.h / utils.cpp图像预处理、仿射变换、可视化等工具函数README.md工程说明文档包含编译步骤和依赖项这个拆分逻辑很实用检测器和姿态估计器是两个独立模型各自封装成独立模块后续替换检测器或姿态模型时互不影响。我在自己的部署项目里也沿用了这个结构把目标检测和下游任务解耦调试起来非常舒服。4.2 RTMDet 检测器的调用人体框预处理与 NMS 后处理RTMDet 是一个高效的实时目标检测器在这里被用来先框出图像中的人体区域。它输出的是解码后的检测框格式为[x1, y1, x2, y2, score]。拿到原始输出后需要做坐标归一化和非极大值抑制NMS。// rtmdet.cpp 中 NMS 后处理的示意逻辑 std::vectorcv::Rect filterAndNMS( float* boxes, int num_boxes, float score_threshold, float iou_threshold) { std::vectorcv::Rect results; std::vectorfloat scores; std::vectorint indices; for (int i 0; i num_boxes; i) { float score boxes[i * 6 4]; if (score score_threshold) { cv::Rect rect( static_castint(boxes[i * 6] / scale_x), static_castint(boxes[i * 6 1] / scale_y), static_castint((boxes[i * 6 2] - boxes[i * 6]) / scale_x), static_castint((boxes[i * 6 3] - boxes[i * 6 1]) / scale_y)); results.push_back(rect); scores.push_back(score); } } cv::dnn::NMSBoxes(results, scores, score_threshold, iou_threshold, indices); std::vectorcv::Rect final_boxes; for (int idx : indices) final_boxes.push_back(results[idx]); return final_boxes; }这里把模型输出的归一化坐标按输入图像与模型输入尺寸的比值还原回原图尺度。scale_x和scale_y需要区分计算因为 TensorRT 输入通常要求固定宽高而原图长宽比不一定一致。直接把整个图 resize 会破坏目标的纵横比常见做法是等比例缩放后 padding项目中也采用了这种方案。score_threshold 一般设 0.5多人密集场景建议降到 0.3避免漏检但会增加后续姿态估计的调用次数。4.3 RTMPose 的后处理SimCC 解码与仿射变换回原图RTMPose 的输出不是坐标而是两个向量。解码的核心是对 SimCC 输出的概率分布取期望。实现方式如下// rtmpose.h 中的 SimCC 解码逻辑 std::vectorcv::Point2f decodeSimCC( float* simcc_x, float* simcc_y, int num_joints, int width, int height, float center_x, float center_y, float scale) { std::vectorcv::Point2f keypoints(num_joints); for (int i 0; i num_joints; i) { // 对水平方向概率分布取期望得到亚像素精度坐标 float expected_x 0.0f, sum_x 0.0f; for (int j 0; j width; j) { float prob exp(simcc_x[i * width j]); expected_x prob * j; sum_x prob; } expected_x / sum_x; float expected_y 0.0f, sum_y 0.0f; for (int j 0; j height; j) { float prob exp(simcc_y[i * height j]); expected_y prob * j; sum_y prob; } expected_y / sum_y; // 将相对模型输入尺寸的坐标映射回原图 keypoints[i].x expected_x / width * scale center_x - scale / 2; keypoints[i].y expected_y / height * scale center_y - scale / 2; } return keypoints; }代码里有个细节注意一开始在 PyTorch 中 SimCC 输出经过了 softmax而 TensorRT 推理拿到的是 logits。所以我这里的实现先做exp()再归一化相当于是手动算了一遍 softmax省下了一个单独的 softmax 层。这样做推理速度更快但代价是如果数值溢出输出会变成 NaN。实际使用中建议对 logits 先做最大值减除数值上更稳定。坐标映射回原图时需要把检测框的中心点和宽高传进来。这里的center_x、center_y是 RTMDet 检测框的中心scale是检测框短边经过 padding 后的边长。如果没有这一步逆向映射所有关键点坐标都会挤在模型输入尺寸的坐标系里画到原图上全是错的。4.4 main.cpp 串联推理流程从读图到可视化整体流程在 main.cpp 中组织成一个标准步骤链读图 → RTMDet 检测 → 对每个人体框做仿射变换 → RTMPose 推断 → 解码关键点 → 可视化结果。int main(int argc, char** argv) { // 初始化检测器和姿态估计器 RTMDet detector(rtmdet.onnx, rtmdet.engine); RTMPose pose(rtmpose.onnx, rtmpose_fp16.engine); cv::Mat image cv::imread(argv[1]); std::vectorcv::Rect boxes detector.detect(image); for (auto box : boxes) { cv::Mat cropped cropAndResize(image, box, 256, 192); std::vectorcv::Point2f keypoints pose.infer(cropped, box); visualize(image, keypoints, box); } cv::imwrite(output.jpg, image); return 0; }这里有一个工程实践中常见的争议为什么不让检测器和姿态估计器共享同一个 TensorRT context我的看法是工程初期两个模型各自独立 context 更安全避免资源竞争导致的隐性崩溃等性能优化阶段再考虑用 stream 方式共享。项目的代码也是这样拆的对新手理解完整流程更友好。5. TensorRT 部署 RTMPose 的常见坑从模型转换到运行时排查5.1 导出 ONNX 后推理结果全零或 NaN现象用 ONNX Runtime 测试导出模型输出的 simcc_x 和 simcc_y 全是 0 或 NaN。原因最常见的是模型加载权重后没有调用model.eval()BatchNorm 层在推理时仍使用训练状态导致数据分布错乱。另一个原因是输入张量的通道顺序错误PyTorch 模型期望 NCHW但测试代码用了 NHWC。解决在导出脚本里强制model.eval()并在导出前用标准化后的输入数据跑一次.pt模型把输出和 ONNX 模型的输出对比。两条输出曲线的数值差异应该在 1e-5 的量级。差异过大的话检查是否有预处理层如 normalize被封装在模型内部如果有导出时就要把对应的均值和方差参数在外部实现避免传递不一致。5.2 TensorRT 构建时报 “Plugin not found” 或算子不支持现象trtexec转换到某个节点时直接报错提示某个插件缺失或算子未被支持。原因ONNX 模型里包含了一些 TensorRT 本身不支持的算子比如GridSample或某些高级上采样变体或者导出时的 opset 版本太新TensorRT 版本解析不了。解决处理方式依次尝试把opset_version降到 13 或 12 重新导出检查模型里是否有自定义算子并考虑用等价算子替换升级 TensorRT 到与 CUDA 匹配的最新版本。项目中带的 ONNX 模型文件已经验证过能直接转换如果你用的是自己重新训练的模型上述排查路径按顺序测即可。5.3 FP16 加速后精度掉点明显现象FP16 引擎推理结果的关键点坐标偏移特别是在小目标或遮挡场景下AP 比 FP32 掉了 3 个点以上。原因SimCC 这个任务的输出是分类向量FP16 的表示精度在处理概率很小但位置敏感的坐标时存在量化误差。尤其是当输出向量的长度较长simcc_split_ratio2.0 意味着输出维度是输入分辨率的两倍尾部的概率分布容易被截断。解决建议对 SimCC 输出层单独保留 FP32 精度其余层使用 FP16。在 TensorRT 构建器 API 中可以通过设置layer-precision实现// 对 SimCC 最后一层强制使用 FP32 auto layer network-getLayer(layer_index); layer-setPrecision(nvinfer1::DataType::kFLOAT);如果项目用的是 trtexec则无法直接指定单层精度——需要改成在 C 里面用TensorRT API方式构建网络并设置层精度。把关键层设为 FP32 后我在自己的项目里实测精度基本能恢复到 FP32 的水平推理速度只损失 5%~10%在可接受范围内。5.4 多路视频输入时显存溢出或帧率骤降现象单路推理正常但扩展到 4 路或 8 路视频流时显存占用急剧上升甚至直接 OOM。原因每路视频流都创建了独立的 TensorRT context 和 CUDA buffer或者引擎的workspace设置过大多路叠加后显存不够用。解决正确的做法是共享同一个引擎engine为每路视频流创建独立的 execution context并用cudaStreamCreate创建独立的流。CUDA buffer 按最大 batch 分配一次各路输入的图像先拷贝到统一的 buffer 池里。项目里没有直接给出多路并发的代码但这个改动路径在 inference.cpp 的封装结构下比较容易改——把 engine 保存为全局单例context 按线程创建即可。如果显存仍然不足优先把输入图像分辨率降到 256×192 以下这个尺寸通常已经能满足绝大多数姿态估计场景的精度要求。6. 性能验证与调参技巧用数字说话拿到能跑的推理工程后下一步就是量化它到底跑多快、准不准。这里的验证流程建议按两步走先做端到端延迟测试再做精度对比实验。延迟测试用 CUDA Event 计时注意要在推理完成后调用cudaDeviceSynchronize()否则计时会漏掉 GPU 上的实际执行时间得出来的数字虚低。我在项目里看 main.cpp 的测量逻辑也是先做多次 warm-up 之后再计时这个习惯值得保持——第一次推理会触发 lazy loading延迟可能是正常值的 3~5 倍。精度对比的标准做法是用同一组测试图片分别跑 FP32 引擎和 FP16 引擎计算关键点坐标的欧氏距离PCK 指标。两者差距在 1.5 像素以内说明 FP16 可以放心用。如果大于这个值就按之前说的方案把 SimCC 输出层改为 FP32。一个额外的调参技巧TensorRT 的optShapes不是设得越大越好。我把 batch 从 1 改到 4 测试后发现延迟反而增加了 20%这是因为内核优化针对的 batch 不匹配实际运行时。正确做法是用你最常出现的 batch size 构建引擎然后用--shapes在推理阶段强行测试其他 batch 的性能曲线找到拐点再做取舍。项目里的动态 batch 支持已经写好了把它利用起来。从那以后我每次拿到新的部署项目都强制自己按这套流程走一遍ONNX 导出后先做数值对比再转 TensorRT 验证算子兼容性最后分别测 FP32 和 FP16 的延迟与精度。这套流程帮我避开了绝大多数部署翻车场景。希望这份工程拆解对你有用如果手上有类似姿态估计的部署需求照着这个项目结构改动会比从零搭建稳妥得多。本文还有配套的精品资源点击获取
返回列表