ARTICLE DETAIL

资讯详情

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

RTMPose实时推理部署:TensorRT优化与C++工程实现

RTMPose实时推理部署:TensorRT优化与C++工程实现 简介针对实时人体姿态估计的算法部署需求这份基于TensorRT优化RTMPose的实战项目面向计算机视觉与深度学习算法工程师与部署人员旨在解决模型在NVIDIA GPU上推理速度慢、资源占用高的问题。项目从RTMPose工作原理与模型结构讲起深入分析TensorRT的层融合、精度校准、内核自动调优和多流并行等优化机制再完整演示模型导出、转换、量化校准、序列化以及反序列化推理流程。实验环节对优化前后的推理速度、准确率和资源消耗进行对比并配合实际案例给出环境搭建、模型加载和实时推理测试方法。压缩包共16个文件以C源码cpp/h、TensorRT engine模型文件为主辅以Python脚本、Visual Studio解决方案sln/vcxproj与README说明整体大小63.6MB目录结构清晰便于按模块理解与复现。已有239人学习/下载适合希望快速掌握高效部署复杂模型、提升算法落地能力的开发者参考。1. 从模型权重到 GPU 毫秒级推理RTMPose 部署到底卡在哪一步很多人拿到一个训练好的 RTMPose 模型第一反应是直接调 ONNX Runtime 或者 PyTorch 跑推理结果发现直播流场景下帧率根本扛不住。问题不在于模型本身而在于你没有针对 GPU 做计算图级别的优化。TensorRT 做的事不是简单地把权重搬进显存而是把网络中的卷积、归一化、激活函数做层融合把精度从 FP32 压到 FP16 甚至 INT8再通过 kernel auto-tuning 为当前显卡挑选最优实现。RTMPose 这种基于 Tokens-to-Token Vision Transformer 结构的模型算子类型比传统 CNN 更复杂部署时对引擎版本、动态 shape 配置、后处理对齐的要求也更高。这个项目给的是一套 Visual Studio 下的 C 工程把 RTMPose 从 PyTorch 导出到 ONNX再转成 TensorRT engine最后在 RTMPose 和 RTMDet 两个模型之间完成了完整的推理链路——前者做人脸或全身关键点后者做人检测两者串联起来才是真正的端到端姿态估计。适合已经在用 TensorRT 但想参考别人怎么处理动态尺寸、多模型串联和预处理对齐的工程师。2. TensorRT 优化 RTMPose 的核心机制与模型转换链路2.1 为什么 CNN 的优化经验不能直接套到 RTMPose 上传统部署流程里ResNet 这类网络结构规整TensorRT 的层融合能轻松把 ConvBNReLU 合并成一个 CUDA kernel。但 RTMPose 的 backbone 是基于 Transformer 的核心操作是 Multi-Head Self-AttentionMHSA里面涉及大量的 reshape、transpose、矩阵乘和 softmax。这些算子在 ONNX 里表达得很碎如果不做图优化每个小算子都会单独 launch 一个 kernel显存带宽和 launch 开销直接吃掉性能。TensorRT 对 Transformer 结构的处理策略主要有两点一是把 QKV 三个线性变换合并成一个大的 GEMM减少 kernel 数量二是对 transpose 和 reshape 这类纯数据搬运算子做消除或合并尽量让数据在寄存器或共享内存里完成流转。RTMPose 的 head 部分还会用到 SimCCSimulated Coordinate Classification来做关键点定位输出的是每个坐标位置的类别概率这部分的 post-processing 通常需要自定义 plugin 或者回到 CPU 做 argmax不能直接依赖 TensorRT 原生算子。2.2.1 导出 ONNX 时最容易埋的坑动态轴与 opset 版本RTMPose 官方仓库提供了导出脚本但如果你直接用默认参数得到的 ONNX 往往是静态 shape输入维度被固定成了类似[1,3,256,192]。这在单张图测试时没问题但一旦你要跑视频流检测框的大小是变化的裁剪出来的人体区域尺寸不可能每次都一样。所以导出时必须显式指定动态轴。import torch from rtmpose import RTMPose # 假设你已经加载了训练好的模型 model RTMPose.from_pretrained(rtmpose-t) # 以 tiny 版本为例 model.eval() dummy_input torch.randn(1, 3, 256, 192).cuda() dynamic_axes { input: {0: batch, 2: height, 3: width}, output: {0: batch} } torch.onnx.export( model, dummy_input, rtmpose.onnx, input_names[input], output_names[output], dynamic_axesdynamic_axes, opset_version17, do_constant_foldingTrue )这段代码里dynamic_axes是关键它告诉 ONNX exporter 哪些维度是动态的。height和width都设为动态是因为 RTMDet 检测出的人体框宽高比不固定如果只动态 batch后面转 TensorRT 时遇到不同分辨率的输入会直接报错。opset_version建议不低于 17太低的话一些 Transformer 里的算子比如aten::mul和aten::add的融合形式可能无法正常导出。导出后先用onnx.checker.check_model校验一遍结构完整性再拿 onnxruntime 跑一次输出确认和 PyTorch 原始输出的精度差距在可接受范围内。2.2.2 ONNX 转 TensorRT 时的网络结构差异拿到 ONNX 之后用trtexec转 engine 是最快的验证方式但如果你想在 C 工程里动态构建就需要用到 ONNX Parser。RTMPose 的 ONNX 里通常包含Einsum或ReduceSum这类算子TensorRT 的 parser 不一定能全部原生支持。我的做法是先跑一遍trtexec --onnxrtmpose.onnx --saveEnginertmpose.engine --fp16看日志里有没有 warning 提示某个算子回退到了 CPU 实现。如果发现Einsum不支持常见做法是在导出 ONNX 前用torch.onnx.export的custom_ops机制或者直接在 PyTorch 里把 MHSA 的矩阵乘展开成显式的matmul和transpose序列。展开之后 ONNX 文件会变大算子数量变多但 TensorRT 更容易做优化。实际测试中展开后的性能反而比直接用 Einsum 高 5% 到 10%因为 TensorRT 能把展开后的 GEMM 序列融合进同一个 CUDA graph。3. C 工程结构与 RTMDet 检测模型的 TensorRT 封装3.1 解决方案里的模块划分项目里的rtmpose_tensorrt.sln是 Visual Studio 解决方案文件核心代码集中在rtmdet.h/cpp、rtmpose.h/cpp、inference.h/cpp和utils.h/cpp这几个文件里。inference.h定义了一个抽象基类提供loadEngine、infer、postProcess三个纯虚函数RTMDet和RTMPose两个类分别继承并实现各自的预处理、推理和后处理逻辑。从工程角度看这种设计的价值在于RTMDet 输出的是检测框坐标和类别分数RTMPose 输出的是 17 个关键点的坐标和置信度两者的网络结构、输入尺寸、后处理算法完全不同。如果写在一个类里后续要替换检测器或者换成别的姿态估计模型时改动范围会很大。拆成两个独立类之后main.cpp 里只需要做对象组合——检测器输出框姿态模型接收裁剪后的图像块。3.2 RTMDet 的预处理细节letterbox 为什么不直接 resize检测模型的输入尺寸通常是固定的比如 640x640。但视频帧是 1920x1080 或者 1280x720直接 resize 会把人体拉变形检测框的坐标回归就会偏。RTMDet 的部署实现用的是 letterbox即在保持宽高比的前提下把图像缩放到能放进 640x640 的最大尺寸剩余区域用 114 填充。// utils.cpp 中的 letterbox 实现省略了边界处理 cv::Mat letterbox(const cv::Mat src, int target_w, int target_h, float scale, int pad_w, int pad_h) { int img_w src.cols; int img_h src.rows; scale std::min(static_castfloat(target_w) / img_w, static_castfloat(target_h) / img_h); int new_w static_castint(img_w * scale); int new_h static_castint(img_h * scale); pad_w (target_w - new_w) / 2; pad_h (target_h - new_h) / 2; cv::Mat resized; cv::resize(src, resized, cv::Size(new_w, new_h), 0, 0, cv::INTER_LINEAR); cv::Mat canvas(target_w, target_h, CV_8UC3, cv::Scalar(114, 114, 114)); resized.copyTo(canvas(cv::Rect(pad_w, pad_h, new_w, new_h))); return canvas; }scale和pad参数必须传递到后处理阶段因为检测框坐标是在 letterbox 之后的图上回归出来的要映射回原始图像尺寸就要先减去pad再除以scale。很多调试半天对不齐框的问题99% 是这里忘了反算。RTMDet 的后处理里还有一层 NMS用的置信度阈值默认是 0.5IOU 阈值 0.45在rtmdet.cpp的postProcess函数里通过nmsThreshold_和confThreshold_两个成员变量控制直播流场景下如果假检多可以把置信度调到 0.6代价是召回率下降。3.3 RTMDet engine 构建时需要注意的 Input 尺寸RTMDet 的 ONNX 导出时输入通常是[1,3,640,640]的静态 shape因为检测模型一般不需要动态输入。但如果你要处理 4K 视频流固定 640 的输入会导致小目标漏检严重。这个项目里 RTMDet 的输入尺寸是在rtmdet.cpp构造函数里写死的想要换 1280 的输入需要重新导出 ONNX 并重新构建 engine。RTMDet::RTMDet(const std::string enginePath, int input_w, int input_h) : inputW_(input_w), inputH_(input_h) { // 读取 engine 文件并创建 execution context std::ifstream file(enginePath, std::ios::binary); // 反序列化 ICudaEngine // 绑定输入输出 buffer } float* RTMDet::preProcess(const cv::Mat img, int pad_w, int pad_h) { cv::Mat letterboxed letterbox(img, inputW_, inputH_, scale_, pad_w, pad_h); // BGR - RGBHWC - CHW归一化到 [0,1] // 返回 GPU 显存上的 float 数组 }注意预处理里的 BGR 转 RGB 和归一化操作TensorRT 的输入是 NCHW 布局的 float 数组不要直接拿 OpenCV 的 Mat 往 GPU 里塞。RGB 转换可以在 GPU 上通过 kernel 做但这个项目是在 CPU 端完成的——输入图是 640x640CPU 转换的开销在 2ms 左右比起整个推理时间占比不小。如果帧率卡在预处理上可以考虑用 CUDA 的cvtColor函数把这一步挪到 GPU。4. RTMPose 的推理实现动态尺寸处理与 SimCC 后处理4.1 动态输入 shape 在代码里的实际处理和 RTMDet 固定 640x640 不同RTMPose 的输入是根据检测框裁剪出来的图像块尺寸变化范围很大。假设检测框是 200x300直接把它 resize 到 256x192会引入畸变。RTMPose 官方训练时用的输入是 256x192但推理阶段更好的做法是保持检测框的宽高比短边 resize 到 192长边按比例缩放然后再 pad 到 256x192。TensorRT 里动态 shape 的处理方式是设置 optimization profile指定每个动态维度的最小、最优和最大尺寸。// 动态 shape 配置对应 inference.cpp 中的 createEngine 逻辑 nvinfer1::IOptimizationProfile* profile builder-createOptimizationProfile(); profile-setDimensions(input, nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, 3, 96, 72)); profile-setDimensions(input, nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, 3, 256, 192)); profile-setDimensions(input, nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(1, 3, 384, 288)); config-addOptimizationProfile(profile);最小尺寸 96x72 对应的是小目标最大 384x288 对应大尺寸人体框。kOPT是最优尺寸TensorRT 会以这个尺寸为准做 kernel 选择和显存分配所以尽量选实际部署时最常见的输入分辨率。我这里把最优设为 256x192和 RTMPose 训练时的输入一致。运行时每次推理前调用context-setBindingDimensions设置当前帧的实际尺寸引擎会自动切换到合适的 CUDA kernel。动态 shape 的代价是显存占用会按最大尺寸预分配如果你的 GPU 只有 4GB 显存把最大尺寸设太大会导致创建 engine 时直接 OOM。一个折中方案是限制最大尺寸到 320x240超过这个范围的检测框直接等比缩到边界内精度损失在小分辨率模型上几乎看不出来。4.2.1 SimCC 输出的坐标解码逻辑RTMPose 的 head 输出不是直接的关键点坐标而是两个方向的概率分布。假设关键点数量是 17SimCC 会为每个关键点生成 x 方向的 192 个类别概率和 y 方向的 256 个类别概率具体数值取决于输出分辨率。解码时先对概率分布做 softmax然后计算期望值再乘上缩放因子得到原始坐标。// rtmpose.cpp 中的关键点解码简化版本 void RTMPose::decodeSimCC(const float* output, int num_points, float* keypoints, float center_x, float center_y, float scale_x, float scale_y) { int W 192; // x 方向的类别数 int H 256; // y 方向的类别数 for (int i 0; i num_points; i) { const float* x_logits output i * 2 * W; const float* y_logits output i * 2 * W W; float x_expect 0.0f; float x_sum 0.0f; for (int j 0; j W; j) { float prob std::exp(x_logits[j]); x_expect j * prob; x_sum prob; } float y_expect 0.0f; float y_sum 0.0f; for (int j 0; j H; j) { float prob std::exp(y_logits[j]); y_expect j * prob; y_sum prob; } keypoints[i * 2] (x_expect / x_sum) * scale_x center_x; keypoints[i * 2 1] (y_expect / y_sum) * scale_y center_y; } }这里为什么不用argmax而用期望因为 argmax 得到的是离散的位置量化误差最大能到半个像素。SimCC 的核心思想就是通过 softmax 期望来做亚像素级别的定位你和 argmax 版本对比一下在手腕、脚踝这些小关节上能差出 3 到 5 个像素这个精度差距在多人场景下直接决定骨架是否抖动。scale_x和scale_y是把模型输出坐标映射回原始图像尺寸的比例因子和 RTMDet 的 letterbox 反算类似。4.2.2 后处理里的 softmax 数值稳定性上面的代码直接对 logits 做std::exp如果某个类别分数特别大exp会溢出。TensorRT 的输出是 FP32RTMPose 经过训练后 logits 的数值范围通常在 -10 到 10 之间直接 exp 没问题但如果换成 FP16 engine数值范围变小溢出概率增加。安全做法是先减最大值再做 softmax但这样会增加一次遍历。实际工程里我通常把输出先拷贝到 CPU用std::max_element找到最大值再做 exp 和归一化。这个操作对性能的影响可以忽略因为输出数据量很小——17 个关键点乘以 448 个类别总共也就 7616 个 float一次 memcpy 加两次遍历的开销在微秒级别。4.3 模型串联时的显存复用RTMDet 和 RTMPose 两个模型不能同时在 GPU 上跑但它们的显存可以复用。TensorRT 的ICudaEngine在创建 execution context 时会分配工作空间如果你的显存吃紧可以在第一个模型推理完成后调用context-destroy()释放显存再创建第二个模型的 context。这个项目里的inference.h抽象基类设计得很好每个模型对象持有自己的 context在RTMPose的构造函数里把RTMDet的 context 传进来先跑检测再释放。5. 推理性能验证与常见工程问题排查5.1 用 chrono 做端到端计时别只看模型推理时间很多人在评估部署性能时只统计 TensorRT 的enqueueV2耗时但这个数字掩盖了预处理、H2D 拷贝、D2H 拷贝和后处理的时间。在 CPU 端做 letterbox 和归一化一张 1920x1080 的图像要花 8 到 12ms如果关键点后处理再写得很粗糙总延迟可能比模型推理还高。这个项目里没有单独写一个 profiling 工具但你可以在 main.cpp 的循环里用std::chrono::high_resolution_clock分段计时。我一般把流程拆成四段读帧预处理、检测推理、裁剪姿态预处理、姿态推理后处理。auto t0 std::chrono::high_resolution_clock::now(); cv::Mat frame capture.read(); // ... auto t1 std::chrono::high_resolution_clock::now(); detector-infer(frame, boxes); // ... auto t2 std::chrono::high_resolution_clock::now(); // 根据 boxes 裁剪并 resize pose-infer(cropped_img, keypoints); // ... auto t3 std::chrono::high_resolution_clock::now();用这种方式跑一段 30 秒的视频统计每段的平均耗时瓶颈一目了然。实际测下来 RTMDet 在 640x640 输入下 FP16 推理约 3msRTMPose 在 256x192 输入下约 2ms但 CPU 端预处理 8ms 成了最大瓶颈。解决思路是改用 GPU 上的cudaMemcpy2D和 CUDA kernel 做 resize、颜色转换和归一化把预处理压到 1ms 以内。5.2 常见问题engine 在不同 GPU 上加载失败TensorRT 生成的 engine 文件是和 CUDA 版本、TensorRT 版本、GPU 架构强绑定的。你在 RTX 3090 上构建的 engine拿到 RTX 3060 上直接反序列化会报错。日志里常见的是engine file is generated on an unsupported platform或者invalid binding。解决办法是先保存 ONNX在目标机器上用trtexec重新构建 engine。项目里的README.md提到了这一点——如果你收到的压缩包里有现成的.engine文件先确认它对应的 GPU 型号和 TensorRT 版本。我通常的做法是写一个自动检测脚本启动时检查当前 GPU 的 compute capability如果不匹配就自动走 ONNX 转 engine 的流程。5.3 精度对比FP32 和 FP16 的输出差异FP16 推理对 RTMPose 这类 Transformer 模型的影响比纯 CNN 模型要大。MHSA 里的 softmax 对数值精度敏感FP16 下可能出现个别关键点偏移 1 到 2 个像素的情况。项目里没有直接对比数据但你可以用nvinfer1::IPluginV2的setOutputType针对特定层保持 FP32。RTMPose 的 head 部分尤其是最后的 SimCC 分类层建议强制 FP32。// 在构建 engine 时对指定层设置 FP32 精度 config-setFlag(nvinfer1::BuilderFlag::kFP16); // 对名为 head.simcc_x 的层设置 FP32 config-setLayerPrecision(network-getLayer(network-getLayerIndex(head.simcc_x)), nvinfer1::DataType::kFLOAT);混合精度配置后engine 文件体积会稍微增大但推理速度几乎不变因为 head 的计算量在整个网络中占比很低。如果你要部署到 Orin 这类边缘设备显存和算力有限这一步尤其值得做——FP16 下精度损失在边缘设备上会被放大因为设备端的 TensorRT kernel 优化策略和桌面 GPU 不同。5.4 多人场景的坑最大检测人数与动态 batch项目里 main.cpp 的实现是一个 batch 只处理一个人体框即 RTMPose 每次只跑一个[1,3,256,192]的输入。假设一帧图像里有 5 个人RTMDet 输出 5 个框代码就要循环跑 5 次 RTMPose 推理。如果每个人裁剪后都重新分配显存、重新 setBindingDimensions总耗时会被放大 30% 以上。优化做法是把多个检测框拼成一个 batch一次性跑 RTMPose。前面动态 shape 配置里的batch维度设为动态运行时把 5 个框 resize 到相同尺寸后 concat 成一个[5,3,256,192]的 tensor推理完成后按 batch 索引取结果。需要注意不同框的宽高比不同统一 resize 到 256x192 会引入畸变折中方案是所有框都 pad 到统一分辨率这样 batch 里的每张图都是同一个网络输入分布。6. 部署到 Orin 等嵌入式平台的调优技巧把这套工程从 x86 平台移植到 Jetson Orin 上时最先遇到的是 TensorRT 版本差异。Orin 自带的 JetPack 里 TensorRT 版本通常比 PC 上的低ONNX parser 支持的算子集也不完全一致。如果你在 PC 上用的 TensorRT 8.6在 Orin 上可能只有 8.5解析同一个 ONNX 文件时警告信息会不同。常见的处理方案是在 Orin 上重新导出 ONNX把opset_version降到 16 或更低避开高版本才支持的算子。Orin 的 GPU 架构是 Ampere 架构的阉割版L2 缓存和显存带宽都比桌面卡小动态 shape 的 kernel 切换开销更明显。实测在 Orin Nano 上把 RTMPose 的优化 profile 最小值设太大会导致显存碎片化严重。我的建议是尺寸范围不要横跨太多档位比如[1,3,128,96]到[1,3,256,192]覆盖绝大多数人体框即可。另一个嵌入式平台的特有问题是 CPU 和 GPU 之间的数据传输。Orin 上 PCIe 带宽比桌面平台低频繁的 D2H 拷贝会把姿态估计的帧率从 30FPS 拉到 10FPS。如果项目后续要接 CSI 摄像头可以考虑用nvarguscamerasrc直接输出 NV12 格式到 GPU 显存省去 CPU 端的颜色转换和 H2D 拷贝。代码层面只需要把cv::cvtColor换成 CUDA 的cudaMemcpy2D加上NvBufferTransform预处理流程的整体延迟能降一半。系统裁剪方面如果你在 Orin 上用 YOLO 系列检测器加 RTMPose 的组合建议把无关的服务停掉比如图形化桌面和 USB 热插拔检测。实时推理任务对 latency 抖动很敏感后台的定时任务或日志刷新会导致掉帧。用ethtool把网络中断绑定到指定 CPU 核心然后用taskset把推理进程锁到另外的核心上实测可以把 99 分位的帧延迟从 180ms 降到 90ms 以下。最后说一个容易被忽略的参数kOPT尺寸里面的数值不要取中间值就完事它直接影响 TensorRT 选择的卷积算法。同样的 RTMPose 模型kOPT设成 256x192 和 224x160生成的 engine 在 Orin 上的推理延迟可能差 15%。所以确定最优尺寸前先统计一下你的应用场景里检测框尺寸的分布取出现频率最高的那个区间而不是直接沿用训练时的输入分辨率。本文还有配套的精品资源点击获取
返回列表