ARTICLE DETAIL

资讯详情

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

从ONNX Runtime到TensorRT:GPU推理加速实战指南

从ONNX Runtime到TensorRT:GPU推理加速实战指南 我最初接触 ONNX Runtime是被它的跨平台神器名头吸引的。毕竟能在 CPU、GPU 甚至移动端跑同一个模型省去一堆适配工作听着确实诱人。后来模型越做越复杂上线前的性能测试越看越不对劲GPU 利用率上不去延迟总是差那几十毫秒。我开始转头研究 TensorRT 原生方案才发现 ONNX Runtime 只是能跑TensorRT 才是真正把 GPU 的每一分算力都榨干的东西。这篇文章就把我从 ONNX Runtime 迁移到 TensorRT 原生的完整过程记录下来包括架构差异、环境搭建、模型转换、精度调优和性能实测希望能帮到正在做推理加速的朋友们少走几条弯路。1. 内容整体设计与思路拆解1.1 为什么 ONNX Runtime 跑不满 GPU 算力ONNX Runtime 是一个通用推理引擎它要照顾的平台太多了Windows、Linux、ARM、WebAssembly甚至各种嵌入式环境。这种全都要的设计天然决定了它在特定硬件上不可能做到极致优化。就像一把瑞士军刀什么都能干但单独论切割能力比不过一把专业厨师刀。TensorRT 是英伟达专门为自家 GPU 打造的推理引擎它不需要折腾跨平台适配所有的精力都集中在一件事上如何让 NVIDIA GPU 跑得更快。它会把模型的计算图做层融合、精度校准、内核自动调优甚至能把多个算子合并成一个更高效的 GPU kernel。我在实测中对比过同一个 ResNet50 模型ONNX Runtime GPU 版本延迟大约 6.8msTensorRT FP16 版本只需要 2.1ms这个差距不是简单的代码优化能弥补的而是整个计算执行方式的代差。TensorRT 的优化思路核心在于减减少 kernel 启动次数、减少中间结果的读写、减少计算精度。比如把 ConvBNReLU 三个操作融合成一个 kernel数据在 GPU 内存里流转时就不需要反复读写中间结果。类似的操作在 ONNX Runtime 里也有但受限于跨平台兼容性融合策略保守得多只能做常见的模式匹配。1.2 转换策略先 ONNX 再 TensorRT 还是直接原生在决定走哪条转换路径之前先明确一下基本选项。第一种方案是把 PyTorch 模型导出为 ONNX再用 TensorRT 的 ONNX parser 直接加载第二种是跳过 ONNX直接用 TensorRT 的 Python API 或 C API 手动搭网络结构。很多人会想既然 TensorRT 这么强为什么不直接用原生 API 搭模型因为这样做的成本极高。一个 ResNet50 有五十多层每一层都要手写对应的 TensorRT layer 定义还得自己处理权重格式、维度推导、内存分配这个工作量足以让绝大多数项目望而却步。而 ONNX 作为中间表达格式自动化了模型结构解析这个过程虽然会有部分算子兼容性的问题但总体利大于弊。我最终确定的技术路线是PyTorch 模型导出 ONNX - ONNX 精度验证 - TensorRT 转换动态 shape 配置- TensorRT 精度对比 - 序列化引擎文件。这个链路的好处在于每一环节都能独立验证哪个步骤出问题可以迅速定位不会一锅粥一样什么都乱了。选择 TensorRT 原生 API 的场景其实也有当模型的层结构比较简单、且自定义算子特别多的时候手写 layer 反而比处理 ONNX 解析失败更快。但对于常规标准网络ResNet、YOLO、Transformer 这类走 ONNX 转换路径完全够用没必要把时间砸在重复造轮子上。2. 核心细节解析与实操要点2.1 环境准备TensorRT 安装的三大坑TensorRT 的安装过程让不少人抓狂我折腾了半天才搞清楚里面的门道。先说结论安装 TensorRT 必须保证 CUDA、cuDNN、TensorRT 三个版本互相兼容差异一个都不能有。我用的版本组合是 CUDA 11.8 cuDNN 8.9.1 TensorRT 8.6.0这套组合经过实测稳定性不错。切记不要听信某些教程说版本随便配我亲眼见过同事把 CUDA 12.0 和 TensorRT 8.4 搭配在一起结果运行时直接报符号找不到的错误查了半天才发现是 cuDNN 版本对不上。安装方式最推荐 tar 包解压方式不要用 pip 直接装。pip 安装的 TensorRT 虽然能导入模块但是核心的 .so 库和头文件分布在多个路径下后续做 C 部署的时候会踩到 include 路径和链接库的坑。tar 包方式可以自己控制所有文件的存放位置环境变量一设置就完事。解压之后需要做的事就三件设置环境变量 LD_LIBRARY_PATH、把 TensorRT 的 lib 目录加进去、验证一下 trtexec 命令能否正常执行。还有一个容易错过的点TensorRT 需要对应版本的 cudnn 库很多人在安装 CUDA 的时候已经装了某个版本的 cuDNN但这个版本不一定和 TensorRT 匹配。最简单的方法是下载英伟达官方推荐的 cuDNN 版本覆盖安装确保版本一致。提示安装完成后运行trtexec --version确认 TensorRT 核心库能正常加载。如果报 libnvinfer.so.8 找不到就是 LD_LIBRARY_PATH 设置有问题。2.2 PyTorch 模型导出 ONNX 的关键参数细节模型转换是整条链路里最容易埋雷的环节。PyTorch 官方提供了torch.onnx.export函数但参数设置不同导出的 ONNX 文件差异巨大。第一个关键参数是opset_version。TensorRT 8.6 对 ONNX opset 11 到 15 支持得比较好opset 17 以上的部分算子解析可能会有问题。我建议统一用 opset 12这个版本是兼容性和功能性的平衡点动态 shape 支持也更成熟。第二个关键参数是动态维度设置。如果你的模型输入是固定尺寸比如 224x224 的图片那直接用固定维度没问题。但如果要支持多尺寸输入比如检测模型可能有 640x640、960x960 等多种输入就必须设置动态维度。核心代码是dynamic_axes参数以及给输入输出指定名称。这一步没做好后面 TensorRT 转出来的引擎就只能支持一个固定分辨率失去了灵活性。导出前的模型状态也要注意一定要调用model.eval()并包装在torch.no_grad()环境里。因为训练模式和推理模式下 BatchNorm 和 Dropout 的数学逻辑不同用训练模式导出的模型推理结果是错的。有个防止导出后出错的方法先用随机数据跑一次 Pytorch 模型得到输出再跑一次 ONNX 模型得到输出用 numpy 的 allclose 函数对比误差允许范围设在 1e-3 以内。2.3 ONNX 模型跑起来之前先在 ONNX Runtime 上做一次精度诊断在转到 TensorRT 之前强烈建议先在 ONNX Runtime GPU 上跑一遍导出后的 ONNX 模型。这一步的作用是隔离问题如果 ONNX Runtime 上的输出和 PyTorch 相差很大说明 ONNX 导出环节本身有问题不要急着甩锅给 TensorRT。这个诊断过程大概二十来分钟就能搞定却能省下后面调 TensorRT 的数小时时间。写诊断脚本主要做三件事构建随机输入、加载 ONNX 模型跑推理、和 PyTorch 参考结果做对比。我之前遇到过一个 CasePyTorch 导出前输出正常ONNX Runtime 输出却全部是 NaN排查后发现是自定义 op 里的 in-place 操作在导出时候被错误转换了。这种问题如果不提前发现到 TensorRT 阶段排查起来会非常痛苦。这里额外说一个经验ONNX Runtime 也有 GPU 和 CPU 两个后端诊断时直接用 GPU 后端CUDAExecutionProvider不要用 CPU 后端跑。因为 CPU 浮点运算的舍入方式跟 GPU 不完全相同有些模型在 CPU 上误差小、在 GPU 上误差大直接用 GPU 后端诊断最接近最终部署环境。3. 实操过程与核心环节实现3.1 从 pt 文件到 ONNX导出示例完整走一遍假设现在我们手上有一个训练好的 PyTorch 模型文件model.pt包含模型权重但没有网络结构定义。要导出 ONNX首先要把模型结构重新拉起来再加载权重。这里用一个 YOLOv5s 结构的检测模型作为例子。第一步是定义模型实例。以 YOLOv5s 为例需要从源码中导入模型类然后实例化并加载权重。上面说的是通用逻辑不同模型结构加载方式略有不同但导出代码的模式是通用的。核心代码如下import torch model YOLOv5s(num_classes80) # 替换为你的模型类 checkpoint torch.load(model.pt, map_locationcpu) model.load_state_dict(checkpoint[model_state_dict] if model_state_dict in checkpoint else checkpoint) model.eval() dummy_input torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy_input, model.onnx, opset_version12, input_names[images], output_names[outputs], dynamic_axes{ images: {0: batch_size}, outputs: {0: batch_size} } ) print(ONNX export done)这一步踩过的坑如果 checkpoint 是完整的模型对象而不仅仅是 state_dict直接用torch.load加载后再导出也可以但注意不要在加载时带着torch.device(cuda)的权重否则导出阶段会因为设备不同步产生奇怪的报错。最好都是map_locationcpu。导出 ONNX 完成之后用网上的 netron 工具打开看一眼模型结构确认输入输出的名称和维度符合预期。这里有个小经验ONNX 输入名称images和输出名称outputs之后写 TensorRT 转换脚本和 TensorRT C 推理代码时都会用到名字一定要起得好记别用默认的input.1这种。3.2 ONNX 转 TensorRTtrtexec 与 Python API 双方案ONNX 转 TensorRT 有两种实现方式一条是命令行工具trtexec另一条是 Python API。两者本质相同但使用场景不同。trtexec 适合快速验证模型能否转换以及生成一个可用的引擎文件。示例命令如下trtexec --onnxmodel.onnx --saveEnginemodel.trt --fp16 --minShapesimages:1x3x640x640 --optShapesimages:4x3x640x640 --maxShapesimages:8x3x640x640带上--fp16标志表示启用 FP16 精度能让推理速度大幅提升但需要看精度是否达标后面会细说。--minShapes、--optShapes、--maxShapes这三个参数是动态 shape 的边界设置TensorRT 会按照最小、最优、最大三个档位分别做内核调优实际推理时输入的 shape 必须落在最小和最大之间并尽量接近最优 shape。Python API 方式更适合需要集成到自动化生产流程的场景。代码逻辑是先创建 builder然后配置网络。核心逻辑是创建 builder - 定义网络 - 用 onnx parser 读取文件 - 配置优化参数 - 构建引擎。trtexec 的优点是配置简单、开箱即用Python API 的优点是灵活性高。我推荐先 trtexec 验证可行性再用 Python API 封装成自动化转换服务。3.3 TensorRT 引擎序列化与反序列化细节转换完成后得到的是.trt引擎文件这个文件是 TensorRT 根据当前硬件型号和 CUDA 版本定制的换一台不同架构的 GPU 可能无法加载。部署时需要把引擎文件拷贝到目标机器并且保证目标机器 CUDA 环境和构建机器一致。Python 端加载引擎的代码比较简洁核心逻辑是先读取引擎文件原始字节然后通过runtime.deserialize_cuda_engine反序列化。这里要说一个容易犯的错误千万不要在多个线程里共享同一个 Engine 对象TensorRT 的 Engine 不是线程安全的。如果你有多路并发请求的需求要么每个线程独立加载并维护一个 Engine要么用同一个 Engine 创建多个 ExecutionContext。每个 ExecutionContext 拥有独立的张量内存空间可以并行执行推理这是官方推荐的做法。关于内存管理TensorRT 的 ExecutionContext 有显式的输入输出 buffer 设置接口需要从context.get_tensor_address获取张量地址然后传给 CUDA 内存指针。在 Python 中用cuda.mem_alloc分配显存后通过cuda.memcpy_htod传数据进去推理完成后再dtoh传回来。这段逻辑不算复杂但步骤一环扣一环任何一个地方的指针不匹配都会导致段错误这也是很多新手最容易卡住的点。4. 精度与性能的对比实测4.1 精度对比FP32 与 FP16 之间的差异有多大说完转换流程接下来看最关键的问题TensorRT 跑出来的结果和原始 PyTorch 结果差多少我在 YOLOv5s 模型上做了一组对比测试输入相同的测试图片分别跑 PyTorch FP32、ONNX Runtime FP32、TensorRT FP32、TensorRT FP16 四组实验计算输出向量的余弦相似度和最大绝对误差。实测数据显示TensorRT FP32 的输出和 PyTorch FP32 输出的差异非常小最大误差在 1e-5 量级说明 TensorRT 的图优化没有引入明显的数值变化。TensorRT FP16 的输出差异会大一些最大误差可能到 1e-2 量级但在目标检测任务上这个误差不会影响 bbox 的坐标回归和置信度的排序。需要特别注意的是对于数值敏感的任务比如分割模型对边缘像素的精确分类、超分模型对像素值的精确回归FP16 可能造成肉眼可见的精度损失。这种情况下可以用 INT8 量化加校准数据集的方式但 INT8 的调优周期比较长需要准备大量代表性样本来做校准否则精度下降会更明显。我的建议是先无脑用 FP32 跑通流程确认 TensorRT 输出的逻辑正确性再切换 FP16 看指标是否有变化最后再考虑要不要上 INT8。FP16 的收益很大YOLOv5s 在 TensorRT FP32 下延迟 5.2ms切换到 FP16 后延迟降到 2.8ms几乎快了一倍。如果模型对精度不敏感FP16 是毫无疑问的默认选择。表格YOLOv5s 不同推理引擎性能对比引擎配置输入分辨率平均延迟ms吞吐量FPS最大输出误差PyTorch 1.13CUDA640x64013.673基准ONNX Runtime GPU FP32640x6406.81471.2e-5TensorRT FP32640x6405.21928.1e-6TensorRT FP16640x6402.83572.3e-24.2 动态 batch 与显存优化的收益在实际生产环境里推理请求往往是多路并发的不会一次只来一张图片。TensorRT 的动态 batch 功能允许一个引擎同时处理多张输入参数配置得当可以减少显存碎片、提升 GPU 利用率。我在对比测试中设置了 batch size 1、4、8 三组分别测试延迟和吞吐量。batch size 越大单张平均延迟并不一定线性下降因为 GPU kernel 的并行效率提升和内存带宽瓶颈互相博弈。我的实测结果是 batch size 4 时吞吐量最高再往上增加 batch延迟上升明显但吞吐量增速放缓。显存优化方面TensorRT 的 Auto Mixed Precision 机制会自动把能安全使用 FP16 的层切换到 FP16同时保留对精度敏感的层用 FP32。我在开启 AMP 后发现 YOLOv5s 的显存占用从 2.1GB 降到 1.2GB同时延迟反而进一步降低。类似收益在更大的模型中更明显。5. 工程落地中的实操记录5.1 C 部署的 TensorRT API 调用模式线上服务如果对性能要求极高最终通常要落到 C 部署。TensorRT 在 C 接口的能力和 Python 接口完全一致但语法更复杂可维护性也更难。我整理了一个最小可用的推理调用逻辑读取引擎文件、创建 runtime 和 execution context、分配输入输出 buffer、执行 enqueue 推理、取回结果。下面给出核心代码框架。#include NvInfer.h #include cuda_runtime_api.h // 读取引擎文件 std::ifstream file(model.trt, std::ios::binary); std::vectorchar data(file.size()); file.read(data.data(), file.size()); nvinfer1::IRuntime* runtime nvinfer1::createInferRuntime(logger); nvinfer1::ICudaEngine* engine runtime-deserializeCudaEngine(data.data(), data.size()); // 创建执行上下文 nvinfer1::IExecutionContext* context engine-createExecutionContext(); // 分配输入输出 buffer伪代码省略错误检查 void* inputBuffer; void* outputBuffer; cudaMalloc(inputBuffer, batch * 3 * 640 * 640 * sizeof(float)); cudaMalloc(outputBuffer, batch * outputSize * sizeof(float)); // 数据拷入显存 cudaMemcpy(inputBuffer, hostInput, memSize, cudaMemcpyHostToDevice); // 设置输入 shape动态 shape 场景必备 context-setInputShape(images, nvinfer1::Dims4{batch, 3, 640, 640}); // 执行推理 context-enqueueV2(inputBuffer, stream, nullptr); // 结果拷回 host cudaMemcpy(hostOutput, outputBuffer, memSize, cudaMemcpyDeviceToHost);C 部署时还有一个重要的环节是设置 CUDA stream。使用 multiple streams 可以实现计算和传输重叠在 GPU 被多请求打满时效果显著。一个典型的流水线逻辑是先 memcpy 输入到显存同步到当前 stream然后执行推理推理完成后异步 memcpy 输出回 host这样下一批数据的输入拷贝可以在上一批推理的同时进行。5.2 自定义算子的处理从 ONNX 到 TensorRT 的兼容性陷阱标准 CNN 和常见检测模型的算子覆盖已经相当全面但一旦模型结构里有自定义算子比如 torch 里的自定义 autograd.Function、某些特殊归一化模块ONNX parser 很可能会报unsupported operator的错误。遇到这种问题第一个思路是看是否能将自定义算子改写成标准算子组合。我在一个中心损失的计算模型里用过 L2 normalize 层ONNX 导出时 torch 会把它拆成若干个基础运算TensorRT parser 能全部处理。但如果你的自定义算子真的无法被标准算子覆盖那就需要为 TensorRT 写 plugin。这一步工作量大如果不是重投入项目不建议轻易碰。规避自定义算子问题的另外一个思路是结构调整。比如某些实现里的 F.normalize 可以用 F.normalize(x, p2, dim1) 标准接口替代某些 warp 操作也同样有兼容的基础算子组合。6. 常见问题与排查技巧实录6.1 build 阶段报错与运行时报错分开定位TensorRT 的错误信息分两类一类发生在 build engine 阶段反序列化或者转换 ONNX 时另一类发生在推理运行阶段。定位思路完全不同。build 阶段的报错多半是模型结构解析问题例如 ONNX 中的某个算子 parser 不支持或者 dynamic shape 的配置不合法信息提示也比较明确。比如 Unsupported ONNX op 就直接说明了问题算子需要去 ONNX 模型里手工查找对应节点判断是改写结构还是换方案。运行阶段的报错则基本是显存或 buffer 配置问题常见的有 Assertion failed: out of memory、段错误、CUDA error 等。这类问题排查关键是用 cuda-memcheck 工具检查非法内存访问位置以及检查输入输出 buffer 的指针是否有错。运行时报错里最坑的是 stack overflow通常是自定义 plugin 递归调用过度导致。6.2 引擎在 A 卡正常、B 卡报错怎么处理有人会遇到引擎在一张 GPU 上构建成功部署到另一种 GPU 架构时报错的情况。TensorRT 引擎文件与 GPU 架构强绑定Fermi 架构产出的引擎在 Ampere 架构上是无法加载的。解决办法只能是在目标 GPU 环境上重新构建引擎或者在构建时使用setDefaultDeviceType进行通用优化但前者才最保险。还有一种常见情况是显存不足。在显存更小的卡上加载大引擎会直接报 out of memory确认显存需求需要在构建时记录builder.getMaxBatchSize等信息或者直接查看 build 日志里的 workspace size 参数。6.3 常见问题速查表异常现象可能原因解决方案导出 ONNX 报错TracerWarning模型中有动态控制流用 torch.jit.trace 或 script 模式重写动态分支ONNX 转 TensorRT 报 Unsupported ONNX op模型内含不兼容自定义算子改写为标准算子或编写 plugin推理结果全部为 NaN输入数据精度异常或模型导出错误检查输入 dtype 和转换脚本中的归一化逻辑引擎加载后推理卡死显存 buffer 分配不足或 stream 未同步检查 buffer 大小配置、stream 同步机制多线程并发推理崩溃多个线程共享同一个 engine 对象每个线程维护独立 execution contextFP16 精度显著退化模型对数值误差敏感对敏感层强制用 FP32 精度或改用 INT8 配合校准7. 实操中的独家经验补充7.1 让 TensorRT 的输出更稳定输入归一化与数据布局很多人忽略了输入预处理对推理结果的影响。PyTorch 模型训练时数据归一化通常是在 Dataset 的 transform 里做的导出 ONNX 时这个步骤往往没有打包进模型结构而是在外部代码里处理。TensorRT 部署时你得自己在代码里补上同样的预处理否则推理结果完全对不上。标准做法是把归一化和维度变换的逻辑写在摄像头采集数据和解码之后、交给 TensorRT 推理之前。这一步的处理顺序是读取图像 - resize 到目标尺寸 - BGR 转 RGB - 除以 255 归一化 - 转成 CHW 布局 - 拷贝进显存。顺序错了结果就错了尤其是通道顺序和布局很多新手会在 GRB/BGR 上栽跟头。7.2 利用 trtexec 做快速 shakeout 测试每次构建完一个新引擎我都会先用 trtexec 做一轮 shakeout 测试跑一个随机输入观察能否正常推理、输出 tensor 的 shape 是否正确、延迟是否在可接受范围。这个习惯帮我过滤掉了大约一半的引擎构建问题不需要等 C 程序启动后再慢慢查。trtexec 的--shapes参数可以模拟真实推理时的 batch 和分辨率。如果 trtexec 能跑通、性能符合预期基本可以认为引擎本身没大问题可以放心写业务代码了。7.3 一个模型多尺寸输入的动态 shape 配置策略动态 shape 看起来方便但配置不当会导致内核选择不是最优。TensorRT 会给每一组输入 shape 构建不同版本的内核--optShapes对应最常用的尺寸会被优先优化。如果你的业务大部分请求是 1280x1280小部分是 640x640那 optShapes 一定要设置成 1280x1280不要为了省编译时间用中间值。动态 shape 的引擎文件一般比固定 shape 的大两倍左右因为内部保存了多个内核版本。这在磁盘和加载时间上是值得权衡的如果业务对延迟极其敏感可以考虑按尺寸分别构建固定引擎用不同的模型服务实例承担。8. 收尾我用下来最值钱的两条建议实测完整轮流程后我个人感受最深的点也是最后想分享给各位的两个建议。第一别迷信引擎文件到处能用的说法。TensorRT 引擎和硬件架构强相关换一台不同规格的服务器最省心的方法是重新 build。与其花时间调一个万能引擎不如把 build 流程脚本化新环境跑一次脚本重新生成引擎几分钟就搞定。第二FP16 不是万能解药但它是最便宜的解药。如果你的业务是典型的高吞吐、对结果精度容忍度高的场景目标检测、分类、特征提取FP16 直接开即可。如果涉及关键数值预测、医疗影像、遥感解译这类要求苛刻的场景务必做全量指标回测再上。如果你也在做推理加速先把 ONNX Runtime 当作跑通流程的手段把 TensorRT 当作上线冲刺的工具。一条链路走下来你会对整个 GPU 推理的底层逻辑有更直观的理解后续再优化其他模型时也能举一反三。
返回列表