最近把一个检测模型从PyTorch搬到TensorRT上做线上推理前后折腾了快两周踩了不少坑也把整个流程彻底跑通了。这篇文章就是想把Ubuntu环境下TensorRT推理部署这套东西从头到尾捋一遍包括环境怎么搭、模型怎么转、C工程怎么写、性能怎么调以及我实际遇到的各类问题和排查思路。适合已经能用PyTorch训练模型、但还没接触过推理加速的工程师也适合部署过但想系统梳理流程的老手。1. 部署方案设计为什么一定要上TensorRT1.1 训练框架做推理的痛点很多团队的第一版上线方案是直接用PyTorch的torchserve或者自写Flask接口加载模型跑推理。这样做的优点是开发快模型不用转换但问题也很明显显存占用高一个batch推理的延迟波动大而且GPU利用率通常只有两到三成。我见过一个分割模型在PyTorch下推理一张图要80多毫秒换成TensorRT FP16之后压到了25毫秒以内差距是数量级的。如果你的模型要跑在实时链路上比如视频流检测、边缘盒子、高并发在线服务继续用PyTorch硬扛会非常吃力。TensorRT本质上是把训练好的模型做一次针对显卡硬件的重新编译和优化。它会遍历算子的多种实现方式组合出在这个显卡上最快的那一条执行路径。这个过程就像同一段代码用不同编译器优化选项重新编译效果差异巨大。PyTorch在推理时依然保留着自动求导、动态图调度这些训练期才需要的机制这些机制在推理阶段纯属额外开销TensorRT把这些全部剔除掉只保留前向推理的最小执行图。1.2 TensorRT的核心加速手段TensorRT的加速不是靠单一技术而是好几层优化叠在一起理解这些才能真正用好它。第一是层融合。显卡执行小算子的开销很大很多算子本身计算量不大但启动和调度成本很高。TensorRT会把相邻的卷积、BN、激活函数融合成一个算子把好几步GPU kernel启动合并成一步。常见的融合方式是ConvBNReLU合并成一个CBR结构实测可以省下不少kernel启动时间。第二是精度校准。TensorRT支持FP16和INT8推理FP16能把计算吞吐提升一倍左右INT8在理想情况下能到四倍。对很多视觉模型来说FP16几乎是无损的INT8需要做校准集和评估但收益明显。第三是内核自动调优。对同一个卷积层TensorRT会尝试多种算法和内存布局选择在该显卡上实测最快的那一个。这个过程就是构建engine时耗时很长的原因自动调优本质上是在跑一个穷举搜索。还有一个经常被忽视的点是显存复用。TensorRT会分析整个计算图的生命周期让不同层的中间结果共用同一块显存大幅降低推理时的显存峰值。这也解释了为什么同一个模型在TensorRT下显存占用可以比PyTorch低不少。我实际测试过一个语义分割模型PyTorch推理时峰值显存约3.2GBTensorRT FP16引擎跑起来稳定在1.1GB左右这个差距在多路并发场景下就是能不能上线的区别。2. Ubuntu环境准备版本选型与安装避坑2.1 基础环境版本怎么选TensorRT对环境的版本匹配非常敏感。CUDA、cuDNN、TensorRT这三个东西的版本必须兼容版本不匹配的典型症状是程序能编译但是运行时报符号找不到或者一跑就core dump。我的建议是优先选一个已经验证过的组合不要全上最新版。TensorRT是跟随CUDA版本走的比如CUDA 11.x配TensorRT 8.x系列是比较成熟稳的梯队CUDA 12.x则适合配TensorRT 10.x系列。选版本时先确定你的显卡驱动支持哪个CUDA注意驱动是向下兼容CUDA的新驱动可以跑旧CUDA反过来不行。然后在显卡驱动支持的范围内往下选一个够用的CUDA版本最后根据CUDA选配套的TensorRT。cuDNN的版本也要对齐TensorRT的依赖要求安装包内通常会有版本说明照着选就行。Ubuntu系统版本我建议用20.04或者22.04这两个版本的glibc和相关库对CUDA生态的支持最稳别用太新的或者太旧的。如果你手头有Docker条件直接拉一个带CUDA的基础镜像再装TensorRT会更省事避开了很多系统库冲突的问题。2.2 安装过程与验证步骤TensorRT的安装推荐用tar包方式这样你能清楚地知道文件装到了哪里排查问题方便。解压之后把lib目录加到LD_LIBRARY_PATH把include目录加到CPLUS_INCLUDE_PATH配置好环境变量。用deb包装会搞得比较黑盒虽然省心但出问题时定位困难不够透明。安装完成之后一定要跑自带的样例验证一遍环境是通的。TensorRT安装目录的samples目录下有很多示例不用全编译只需要验证trtexec这个工具能不能正常跑。或者更简单直接找一张图片换成一个简单的ONNX模型用trtexec转一下能生成engine就说明基础环境没问题。常见的报错我列一下启动时提示找不到libnvinfer.so说明LD_LIBRARY_PATH没配置对或者压根没安装到位提示libcudnn.so找不到说明cuDNN的lib路径没加进去提示CUDA driver version is insufficient说明显卡驱动版本太老需要升级驱动而不是CUDA toolkit版本。这类环境问题占据了部署排障的很大一部分把基础环境弄清爽了后面的模型转换和C编译会顺畅很多。2.3 显存和硬件的现实约束TensorRT部署还存在一个前提约束显卡算力版本太老的话有一些优化用不了。在下载TensorRT安装包或者看官方文档的适用范围时注意一下有没有写清楚支持的算力范围比如对INT8量化一般需要在特定算力版本及以上才支持。如果是比较旧的显卡可以先跑通FP16INT8的收益可能达不到预期。显存方面也有一个容易被忽略的点TensorRT构建engine时有一个工作区上限设置控制自动调优时可使用的显存空间。工作区给得越大越有机会找到更优的kernel实现但也会占住显存。新版TensorRT里这个参数已经改名为memory pool limit但本质上就是那个workspace size。我在构建一个较大的分割模型时试过把工作区从1GB放宽到2GB推理速度提升了大约百分之十五但显存峰值也同步上去了这个取舍要看实际业务场景没有绝对最优值。3. 从PyTorch到ONNX模型转换的关键细节3.1 导出ONNX时最容易踩的坑在把PyTorch模型转成ONNX时很多人第一步就卡住了。torch.onnx.export这个接口看起来简单但导出出来的ONNX文件质量会直接决定后面TensorRT的优化空间。首先要指定合适的opset版本。opset太老很多算子不支持opset太新TensorRT的解析器可能还没完全跟上。我一般固定用opset 11到13这个区间稳定性和算子覆盖率都比较理想。其次要注意动态轴。如果你的模型要在推理时支持可变batch或者可变输入尺寸导出时必须显式指定dynamic_axes否则导出的ONNX会把输入shape写死TensorRT那边再怎么配动态shape都白搭。还有一个不起眼但很烦人的坑模型中如果用了torch.where这类会产生额外算子节点的操作ONNX里可能解析成多个算子组合TensorRT优化时容易卡住或者报不支持。遇到这类情况能重写就重写比如把一些条件逻辑放到预处理里面不要留在模型计算图里。另外一个很重要的建议是导出后要验证ONNX的中间张量形状是否符合预期。可以直接用它跑一次onnxruntime推理对比输出和PyTorch的差异。这一步能提前暴露很多算子转换错误不要等到TensorRT阶段再排查。3.2 用trtexec快速验证转换效果拿到ONNX之后第一件事不要急着写部署代码先用trtexec跑一遍转换和推理基准。trtexec是TensorRT自带的一个命令行工具它能完成从ONNX转换到engine、执行推理、输出性能统计的完整流程。这一条命令能验证的事非常多ONNX解析器是否支持模型中的所有算子、当前显卡上转换是否成功、FP16能不能正常开启、推理延迟和吞吐是多少。实际使用时我一般这样跑trtexec --onnxmodel.onnx --saveEnginemodel.engine --fp16 --workspace2048加上--fp16开关模型会使用FP16精度推理如果转换过程没有报错说明模型结构对FP16是友好的。如果这里就报错那大概率是ONNX里面存在问题算子要在导出阶段修复而不是在TensorRT层面绕。还要注意引擎构建时给workspace配的记忆池大小命令里的单位是MB我给2048表示2GB工作区。trtexec跑完之后输出的指标很有参考价值重点关注latency的p50、p99以及throughput的数值他们代表你这个模型在当前显卡上的理论上限。后面写的C代码如果没有达到这个水平说明工程实现里还有优化空间这个基线数据就是你的对比锚点。3.3 polygraphy精度对比的实用工具trtexec能解决能不能跑通的问题但解决不了精度对不对的问题。模型可能转成功了、跑起来了但输出结果跟PyTorch差距很大这种时候用polygraphy来做精度对比非常方便。polygraphy是一个专门的模型调试工具可以用来对比不同后端、不同精度下的输出差异。我用它做的最多的事情是把PyTorch的输出和TensorRT引擎的输出放在一起对比数值误差定位是哪一层开始出现偏差。操作起来不算复杂指定模型路径和输入数据它会自动跑两个后端并输出逐层或者逐算子的差异分析。在实际项目里我遇到过FP16模式下某个卷积层的误差被放大器逐渐放大模型整体输出完全不能用的情况。用polygraphy逐层对比之后发现问题出在某一层的数值范围差异过大导致FP16精度不够用。最后通过在该层前增加一个归一化节点解决了问题。这个排查过程如果靠肉眼对输出图根本定位不到原因polygraphy的作用就在这里。建议TensorRT部署的流程里把它加进去和trtexec搭配使用一个解决速度问题一个解决精度问题项目会稳很多。4. C推理代码的工程化实现4.1 序列化与反序列化流程用Python写推理代码是最简单的但生产环境要追求极致性能和低延迟还是得用C。TensorRT的C API设计和它的执行模型是一一对应的把核心概念理清楚之后代码写起来并不复杂。整体流程是这样的第一步从ONNX构建engine这一步在线上环境通常不做而是在离线阶段完成然后把engine序列化保存成文件。第二步应用启动时从engine文件反序列化加载得到ICudaEngine对象。第三步创建IExecutionContext执行上下文分配输入输出显存然后循环调用执行接口跑推理。最后在每次推理前后做显存和主机内存的拷贝就行。这里有一个对新手很友好的设计engine文件可以跨进程复用不需要每次启动都重新构建。构建engine是最耗时的环节一个稍大点的模型可能要几分钟甚至更久而序列化后的engine加载只需要几百毫秒。所以在生产环境里的标准做法是离线构建线上加载这是工程化部署中最核心的一环。反序列化时TensorRT要求提供一个Logger对象来接收日志信息。如果日志级别设置不正确可能会漏掉关键的错误信息。我习惯把日志级从WARNING调到INFO虽然在客户环境里太啰嗦但自己开发调试阶段信息越全越好。不需要的日志可以通过重载Logger类的log方法做过滤代码量不大但是排查问题时的体验差别很大。4.2 输入输出管理与显存优化在C API里最容易让人困惑的部分是指针和显存管理。核心要理解的是IExecutionContext的execute接口接受的是绑定指针数组这些指针指向的是GPU显存地址不是CPU内存地址。这意味着推理之前你必须把输入数据从CPU拷到GPU用cudaMemcpyAsync把数据搬到绑定的输入显存上然后执行推理推理结束后再从输出显存拷贝回CPU。这里我建议统一使用cudaMemcpyAsync配合CUDA流而不是用同步的cudaMemcpy因为异步拷贝可以和kernel执行重叠后续做流水线优化时更灵活。为了减少每次推理的分配开销我倾向于在初始化阶段就把输入输出显存一次性分配好之后循环复用。每次推理只是覆盖数据内容不重新分配显存可以有效避免频繁cudaMalloc带来的开销和显存碎片。这里的显存碎片问题是容易被低估的程序长时间跑着频繁分配释放不同大小的显存块最终可能导致显存利用率变低新分配时出现莫名其妙的显存不足。解决办法就是一初始化就分配好固定大小的buffer全程复用。还有一个小细节是数据预处理尽量放到GPU上做。图像缩放、归一化、通道转换这些操作原本如果用CPU做然后再拷到GPU会有一次D2H再H2D的额外开销。如果在一开始就把原始图像拷到GPU显存然后用GPU算子完成预处理整体的端到端延迟可以降低不少。我第一次做TensorRT部署时就把预处理留在CPU端结果发现即使推理本身只要20毫秒加上预处理和两次拷贝整体延迟翻了一倍。后来把预处理迁移到GPU上才把端到端延迟真正压下来。4.3 多路并发与异步推理单路推理优化到一定程度之后性能提升空间就有限了这时可以往多路并发方向走。TensorRT的多路并发有两种常用模式一种是多线程多context每个线程创建自己的ExecutionContext共享同一个engine。这样可以利用多核CPU并行提交任务每条线程有自己独立的执行上下文和绑定的显存互不干扰。另一种是单线程内做流式异步推理适合单路视频流这种需要降低单帧延迟的场景。多线程共享engine时有一个注意点ICudaEngine是可以多线程共享的但IExecutionContext不能同时被多个线程调用每个线程必须有自己的执行上下文。如果多个线程共用一个context会出现随机崩溃或者结果错乱。我在最初设计并发模块时忽略了这一点线上跑了几个小时之后出现了偶发性的结果错误排查了很久才定位到是并发访问context导致的这个教训相当深刻。创建多个context时每个context都会有自己的activation显存和workspace需求要留意总显存是否超限。如果资源比较紧凑可以多用几个较小的context实测下来对单个模型的并发吞吐是有帮助的。另外在新版TensorRT中跨context共享显存池API已经能用但具体策略仍然需要结合实际部署的显存预算来权衡。4.4 关于新版API的兼容提醒TensorRT的大版本更新会引入API变化最明显的是新版执行接口从enqueueV2逐渐演进到了enqueueV3。V3版本引入了更清晰的依赖关系描述和显存分配方式对异步场景设计更友好但代码写法跟老版本差别很大。如果项目是长期维护的我的建议是根据TensorRT版本来选择对应的适配方案不要跨版本硬复用代码。老版本项目里别急着全量升级到新API先把当前版本跑稳定再说新项目可以直接学习新接口但要做好踩坑的心理准备。一般情况下只要在工程入口处做好版本判断把创建context和执行推理的代码集中封装在一个适配层里后续升级TensorRT版本时只需要改一个模块就够了。5. 性能调优与INT8量化实践5.1 先做性能基线再逐项优化性能优化最忌讳凭感觉东改一下西改一下。正确做法是先跑出当前性能基线列出所有影响性能的环节然后逐项测试对比。影响TensorRT推理性能的因素大致可以分成几层模型本身的复杂度、TensorRT引擎构建参数、输入输出拷贝路径、预处理的位置、并发策略、以及与业务代码的交互方式。我建议在调优时先通过trtexec拿到模型在空转条件下的理想延迟再用自己的工程代码测真实端到端延迟二者对比就能知道工程层面还有多少优化空间。如果实际延迟明显高于trtexec结果优先检查数据通路。有没有不必要的主机内存拷贝预处理是不是放在CPU上CUDA流有没有正确使用这些都是性能损失的大头。如果实际延迟已经接近trtexec水平还想继续压那就得从引擎层面找办法比如增大工作区、关掉动态shape的灵活性、或者启用INT8量化。5.2 INT8量化实测与校准集构建INT8量化是TensorRT里收益最大也是最危险的一环。FP16通常是无损的但INT8需要通过校准集统计每层激活值的分布然后把浮点数值映射到INT8的表示范围。校准集的质量直接决定量化之后的精度这一环没有做好模型整体精度可能会崩。校准集要尽量覆盖线上真实场景的输入分布数量不用太多我见过几百张到一两千张图片都能取得不错的效果。关键是多样性别只用一种光照、一种角度、一种目标的图片。构建校准集的过程在TensorRT里是通过实现一个校准器类来完成的它会读取一批输入数据统计激活分布然后计算每个张量的缩放因子。TensorRT支持逐张量per-tensor和逐通道per-channel两种量化粒度。逐通道的精度通常更好尤其是在通道间数值分布差异大的网络中但代价是可能增加一些算子的计算开销。我在部署一个目标检测模型时试过两种方案逐通道的mAP比逐张量高了约0.8个点推理速度几乎没差别。如果精度压力比较大的项目值得优先尝试逐通道。校准完成后需要用验证集做充分的精度评估不能只在训练集上测。量化掉点如果控制在一定范围内属于正常现象但如果超过预期就要考虑是不是校准集分布偏差太大或者某些敏感层不适合INT8需要开启层级精度回退。5.3 容易被忽视的性能陷阱有几个性能陷阱是新手最容易忽略的。第一个是动态shape的副作用。如果引擎支持动态输入尺寸TensorRT在运行时每次可能会重新选择kernel实现而不是走最快的固定路径。如果业务允许尽量在构建引擎时固定输入尺寸或者只提供少量几个预设尺寸这会带来明显的性能提升。第二个是batch太小导致GPU利用率不足。很多推理服务是单请求进来就跑一次batch默认为1GPU的并行能力其实没有发挥出来。如果能做动态batching把多个请求攒在一起推理吞吐量会有数倍提升。我实现过一个简单的动态batching模块把同一时间窗口内的请求合并成batch线上吞吐从每秒钟处理几十次提升到几百次效果显著。第三个是数据拷贝与推理没有重叠。如果使用cudaMemcpy同步拷贝CPU会一直等待GPU完成任务GPU也可能空闲下来等待CPU准备数据两个设备都在各自等待时间就浪费了。正确的做法是用多个CUDA流把数据预取、推理、结果回传三个环节做成流水线让CPU和GPU并行工作。改造之后端到端延迟即便没有下降吞吐能力也会上一个台阶。6. 常见故障排查清单6.1 启动与构建阶段报错这一类错误的特点是发生得早通常在程序启动或者engine构建阶段就会暴露定位路径相对清晰。报错现象可能原因排查思路找不到libnvinfer.soLD_LIBRARY_PATH未配置检查环境变量是否指向TensorRT解压目录的lib文件夹找不到libcudnn.socuDNN库未安装或路径不对确认cuDNN版本与CUDA版本匹配加入ld路径提示CUDA driver version is insufficient显卡驱动过旧升级驱动而不是升级CUDA toolkit构建engine时提示opset不支持ONNX导出时opset版本太高或太低重新导出ONNX固定opset 11到13之间构建时显存不足工作区配置过大或模型过大降低workspace上限或换更大显存机器报错Failed to register plugin模型需要特定插件检查模型是否包含自定义算子确认插件库是否加载如果构建engine阶段出现解析失败的报错信息里通常会指出是哪个节点或者哪种算子出问题。把这一行信息保留下来去PyTorch源码里找到对应层看能不能用等价算子替代往往比在TensorRT层死磕更快。6.2 精度与推理结果异常这一部分的问题最让人头疼因为程序能跑但结果不对定位起来比较费劲。先说一个最常见的场景FP16模式下模型输出和PyTorch结果对比数值偏差明显。这种情况不一定是TensorRT的问题有可能是模型本身的动态数值范围过大FP16的表示精度扛不住。排查思路是用polygraphy逐层对比定位第一层出现明显偏差的位置然后在该层附近增加数值归一化或者改用FP32运行该部分。还有一个场景是INT8量化之后精度大幅下降。可能性最大的是校准集分布与真实数据差异过大或者校准集样本量太少。我遇到过用训练集的子集做校准效果还行但换成线上真实数据后精度崩了的情况。解决的办法是尽可能用线上真实采集的数据构建校准集而且要多场景多时段采样保证分布覆盖。如果输入输出数据的shape不一致导致结果解析错位通常是因为动态shape配置时的命名没有对应上。检查导出ONNX时dynamic_axes的键名和C代码里绑定张量的名字是否一致名字对不上就会拿错数据。6.3 性能不如预期性能不达标的问题建议先回到trtexec的基准数据做对比。如果trtexec跑出来的延迟和吞吐已经很好但你的工程代码达不到重点排查数据通路。具体而言是否用了同步拷贝预处理是否在CPU上是否每次推理都重新分配显存是否单请求单推理没有利用batch能力。还有一种隐性性能损耗来自CPU和GPU之间的同步等待。如果代码里有不必要的cudaDeviceSynchronize比如每次推理后为了拿结果就强制同步一次会打乱GPU流水线。正确做法是使用事件或者流回调来管理完成时机让CPU不用干等GPU执行结束。如果trtexec本身的表现就一般说明问题在模型或者引擎配置。尝试增大工作区、固定输入尺寸、开启FP16或者INT8这些手段的效果需要逐一测试不要一次性全开否则定位不了哪个改动真正起了作用。我自己在性能调优时习惯记录一个对比表格把每次修改前后的延迟、吞吐、显存峰值记下来。这个习惯在优化后期特别有用因为很多参数之间会互相影响如果没有记录改来改去很容易绕晕。做了几年部署相关的工作我最大的体会是TensorRT的性能上限其实很高大多数项目达不到那个上限原因不在TensorRT本身而在工程链路里的数据通路和并发设计。把基础的版本环境理清楚用trtexec和polygraphy建立起速度和精度的基线再把C侧的数据流和并发结构做好性能和稳定性就有了保障。这套流程我后来又套用到好几个不同类型的模型上每次都很稳。如果你们正在做类似的事情希望这篇文章能帮你们少走几段弯路。
到O(n)的双指针原地覆盖)
