资讯动态

TensorRT第一个程序:从ONNX到Engine的模型编译实战

发布时间:2026/10/21 12:11:35 来源:尧图企业网站定制
1. 先说清楚TensorRT到底在解决什么问题1.1 它不是框架是个“编译部署工具”很多刚接触TensorRT的开发者第一反应是TensorRT是不是又一个深度学习框架像PyTorch一样用来训练模型不是。TensorRT是NVIDIA推出的推理优化引擎和运行时库它的工作对象不是训练而是“把已经训练好的模型编译成一串能在GPU上跑得更快的部署程序”。我这么说可能还是有点抽象换个方式理解训练好的模型就像一个“设计图纸”里面画满了算子Conv、ReLU、Pooling、MatMul等。PyTorch或者ONNX Runtime拿到这张图纸是一边看图纸一边按顺序找对应工具去干活这种方式灵活但每道工序都有调度开销。TensorRT拿到图纸后会把很多工序合并成一道把该提前算好的提前算好还能给每条流水线挑最熟练的工人然后把整套流程打包成一个高度定制化的“专用设备”。这个“打包”动作就是我们常说的构建Engine模型编译。所以这张图景就很清晰了平时说的“部署模型”第一步通常是把PyTorch导出成ONNX第二步就是用TensorRT去解析ONNX、优化图结构、生成推理计划第三步拿到编译好的Engine文件.engine/.plan写一个轻量级的推理程序调用它。标题里说的“第一个trt程序实现模型编译的过程”聚焦的就是第二步。你不需要自己手写CUDA算子但要能把这套编译流程跑通并且知道每一步在干什么。1.2 模型编译 vs 直接推理“编译”到底编了什么直接跑ONNX比如用onnxruntime和先编译成TensorRT Engine再跑区别在哪我用一个很实际的项目场景说。之前我给一个分类模型做部署初期直接用ONNX Runtime推理一张图的耗时大概在2ms出头。后来换成TensorRT编译后再推理耗时压到了1.1ms左右且这个差异会随着模型变大、算子变多变得更明显。原因不在于TensorRT能“学会更好的算法”而在于它把大量优化工作放到了编译期提前完成。编译期内TensorRT至少干了这几件事张量内存复用神经网络中间会产生一堆残留特征图如果不做规划每个算子都临时要显存、用完再释放碎片化会吃掉很多空间和调度时间。TensorRT会把整张计算图的内存需求整体做一次静态分析像行李打包一样把生命周期错开的张量放到同一块显存里。这点对生产环境特别重要显存占用经常会降下来30%甚至更多。层融合Kernel Fusion最常见的是把Conv BN ReLU这类固定组合融合成一个算子减少显存读写和内核启动次数。做图像模型部署的朋友应该都听过“ConvBNReLU融合”这个经典操作TensorRT在编译期会自动找这种能合并的模式不用你手工改模型结构。算子选择Tactic selectionGPU上同一个卷积有很多种实现方法比如直接卷积、im2colGEMM、Winograd、FFT等每种在不同shape、不同GPU上表现不一样。TensorRT会在编译期利用一个“策略扫描”的机制在目标卡上实测一批候选实现挑耗时最短的写进Engine。这就是为什么编译时通常很慢但推理时特别快。精度校准如果开FP16或INT8还需要在编译期做精度缩放和校准把原本FP32的权重和激活值转成更低精度同时尽量把精度损失控制在可接受范围。这些动作叠加在一起最后生成的Engine文件已经不是一个传统意义上的“模型参数文件”而是一组针对特定GPU架构、特定输入shape、特定精度写好的可执行推理计划。这也是为什么Engine文件换一张不同计算能力的卡就可能加载失败需要针对目标卡重新编译。2. 环境准备这一步踩的坑比后面所有代码加起来都多2.1 版本匹配第一原则写第一个trt程序之前环境能不能配对基本决定了你后面是顺畅还是折腾。我见过太多人在这一步卡了两三天。TensorRT对版本极其敏感需要同时对齐至少三件事CUDA版本cuDNN版本TensorRT工具包版本它们不是“越大越好”而是配套关系。比如你的GPU驱动支持CUDA 12.x那你装TensorRT 10.x通常没问题如果系统里是CUDA 11.x却硬装TensorRT 10.x大概率会在加载库的时候直接报undefined symbol一类的错误。我的建议是先确定你要用的TensorRT版本去NVIDIA官网看它对应的CUDA、cuDNN最低版本要求然后照着配。不要先随便装个最新CUDA再去碰运气。如果你用的是官方容器或者镜像那更省事基础环境基本已经给你组合好了只需要确认自己的GPU架构能被支持。如果是自己装强烈建议全程记录版本号哪个环节出问题第一件事就是排查版本组合。2.2 两种安装方式与验证方法TensorRT的安装有几种形态我最常用的是tar包解压安装因为能清楚看到目录结构也方便在部署环境里拷贝运行库。核心步骤大概是下载TensorRT tar包解压到某个目录比如/opt/tensorrt。把/opt/tensorrt/lib加入LD_LIBRARY_PATH。Python端的话进入TensorRT的python目录安装对应的whl包例如pip install tensorrt-10.0.1/python/tensorrt-10.0.1-cpXX-cpXX-linux_x86_64.whl验证是否装好python -c import tensorrt; print(tensorrt.__version__)如果能正常打印版本号说明库已经能被找到。这里有个很容易踩的坑whl包里只是Python接口底层还是要动态链接TensorRT的运行库。如果你解压了tar包但没有设置LD_LIBRARY_PATHimport tensorrt 可能报错提示找不到libnvinfer.so。这时候不要怀疑是pip安装坏了多半就是环境变量没指过去。另外如果你想在编译期跑ONNX解析器还需要确认Python的onnx包也装了并且版本不要太老。TensorRT的OnnxParser是基于标准ONNX定义的onnx版本太老会导致解析某些新格式时奇奇怪怪地失败。我自己一般装最新的稳定onnx版本基本不会和TensorRT冲突。3. 第一个trt程序从ONNX到engine的完整实现3.1 一个可以直接抄作业的编译脚本环境准备好之后我们就可以写第一个trt程序了。下面这个脚本是我在自己项目里最常用的一套骨架它负责把一个ONNX文件编译成TensorRT Engine文件并且支持选择FP32或FP16。import tensorrt as trt # Logger是必须的控制TensorRT日志输出 TRT_LOGGER trt.Logger(trt.Logger.WARNING) def build_engine(onnx_file_path, engine_file_path, halfFalse): # 1. 创建 Builder builder trt.Builder(TRT_LOGGER) # 2. 创建网络定义EXPLICIT_BATCH表示显式指定动态/静态batch network builder.create_network( 1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) # 3. 创建 BuilderConfig所有编译配置都挂在它上面 config builder.create_builder_config() # 4. 创建 ONNX Parser把onnx文件解析进network parser trt.OnnxParser(network, TRT_LOGGER) # 5. 读取并解析ONNX with open(onnx_file_path, rb) as f: if not parser.parse(f.read()): print(ONNX解析失败错误信息) for i in range(parser.num_errors): print(parser.get_error(i)) return None print(ONNX解析完成网络输入数量, network.num_inputs) print(ONNX解析完成网络输出数量, network.num_outputs) # 6. 配置常用选项 # 这里给workspace设置1GB编译期显存规划空间 config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) if half: config.set_flag(trt.BuilderFlag.FP16) # 7. 编译返回序列化的engine字节流 serialized_engine builder.build_serialized_network(network, config) if serialized_engine is not None: with open(engine_file_path, wb) as f: f.write(serialized_engine) print(Engine已保存到, engine_file_path) return serialized_engine else: print(编译失败) return None if __name__ __main__: build_engine(model.onnx, model.engine, halfTrue)这段代码是TensorRT 8.x之后、10.x之前的常用写法。如果你是特别老的版本可能会看到builder.build_engine返回一个ICudaEngine对象再通过engine.serialize()拿到字节流但新版本推荐直接用build_serialized_network一步到位返回序列化结果代码更简洁。3.2 逐行拆解几个关键对象上面这段代码看起来不长但背后每个对象都有讲究我一次性讲清楚。Loggertrt.Logger(trt.Logger.WARNING)是TensorRT打印日志用的对象任何TensorRT程序都躲不开它。日志级别有VERBOSE、INFO、WARNING、ERROR。平时日常编译用WARNING就够了输出干净如果排查问题可以临时改成VERBOSE能看到每个算子融合、每个tactic实测的详细信息但编译速度会明显变慢、日志量大到刷屏。Builder、Network、Config这三个对象是铁三角。Builder是总指挥Network描述的是“从ONNX解析出来的原始网络结构”Config记录的是所有编译策略。为什么要把Network和Config分开因为同一个网络结构可以用不同配置编译出不同的Engine比如一个FP32的、一个FP16的、一个INT8的。这种设计让模型结构定义和优化策略解耦我们在迭代精度时非常方便。EXPLICIT_BATCH标志这里必须用1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)这个写法。如果不用它TensorRT默认走隐式batch模式很多动态shape操作会受限。现在是2024年了我建议一律使用显式batch应对固定batch和动态batch都灵活。OnnxParsertrt.OnnxParser(network, TRT_LOGGER)负责把ONNX文件按顺序读出来转成TensorRT的Network定义。它本质上是做一层“翻译”把ONNX里的OP映射成TensorRT支持的Layer。如果遇到TensorRT不支持的算子parser就会报错并记录到parser.num_errors和parser.get_error(i)里通常错误信息会明确指出是哪个节点、哪个算子出了问题。我第一次跑通这个脚本时最大的感受是编译的核心流程其实就这么几个对象API设计得很直接真正让新手头疼的反而是各种配置项的含义。3.3 动态shape配置三行代码决定你的batch灵活性如果你的模型输入shape是固定的上面那段脚本直接用就好。但如果想用一个Engine同时支持不同batch大小或者支持长宽可变的图像输入就必须配置Optimization Profile。Optimization Profile说白了就是告诉Builder我未来可能用哪些shape范围你按这个范围去优化。下面是一个典型的动态batch配置示例假设输入张量名是inputprofile builder.create_optimization_profile() # 这里假设network的第0个输入名是input input_tensor network.get_input(0) # 分别指定 min shape / opt shape / max shape profile.set_shape(input, (1, 3, 224, 224), (4, 3, 224, 224), (16, 3, 224, 224)) config.add_optimization_profile(profile)这三个shape的作用分别是min shape允许的最小shape低于这个范围运行时直接报错opt shape优化基准shapeBuilder会重点按这个shape调优实际推理shape越接近它性能越好max shape允许的最大shape用于显存分配上限规划。需要注意的是profile一次只能覆盖一个shape范围如果有多个差异很大的范围可以创建多个profile并全部添加到config里运行时再按需求选择。但profile越多编译时间和显存占用也越高不要贪多。我一般在生产环境只放一到两个最常用的profile覆盖训练场景和上线场景就够了。4. 模型编译过程中真正影响性能的四个选项4.1 精度选择FP32、FP16、INT8怎么选很多人在第一次编译时就面临选择到底是FP32还是FP16要不要试INT8我的判断标准很简单先看精度要求再看硬件支持最后看显存和延迟瓶颈。FP32最稳和训练精度基本一致适合第一个版本先跑通流程FP16半精度绝大多数现代GPU都支持速度通常比FP32快一倍上下显存占用缩半是部署最主流的默认选项INT8需要校准数据量化后精度会有一定下降但延迟和显存又进一步压缩适合对吞吐有极高要求、精度损失能接受的场景。在BuilderConfig里开启FP16只需要加一行config.set_flag(trt.BuilderFlag.FP16)。但如果你的GPU不支持FP16计算编译时会直接失败或者TensorRT日志里出现明确提示。现在市面上主流N卡基本都支持除非你是特别老的架构。下面这个表格是我在实际部署中比较不同精度的直观感受具体数值和模型有关但趋势是稳定的精度相对耗时显存占用精度风险适用场景FP32基准基准极低首次上线、对比验证FP16约50%~70%约一半一般可忽略绝大多数生产场景INT8约25%~40%可进一步降低需要校准验证高吞吐、视频流处理我个人建议不要一上来就追求INT8。先FP32把整条链路跑顺然后切FP16看精度如果精度还在容忍范围就用FP16。INT8是后面性能死活上不去的时候再考虑的选项。4.2 workspace与tactic为什么编译期那么慢第一次跑编译脚本的人都会问一个问题编译一个模型怎么要几分钟这不像是在“转换模型”更像是在跑一个优化算法。答案是TensorRT要对每个算子做tactic扫描。tactic可以理解为某一层在目标GPU上“候选内核实现策略”不同tactic对应不同的拼接内存方式、计算指令、并行划分。编译时Builder会实际调度这些候选kernel跑一遍基准测试选出最快的那个。这里就涉及到workspace配置。我在脚本里写的是config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30)1 30就是1GB字节。这是指Builder在编译期间和最终Engine运行时能用于临时计算和算法候选的显存上限。给得越大Builder越有可能尝试更多tactic找到更优实现但编译时间和显存压力也会上升给得太小一些耗显存的算子就只能退而求其次性能可能打折扣。所以这是一个平衡问题。我的经验是先给1GB跑一版然后看日志里是否有显存不足的警告如果模型很大可以给2GB或者4GB如果服务器显存紧张就压缩workspace别让它和运行时显存打架。4.3 engine的复用编译一次到处加载生成.engine文件只是第一步。实际部署时我们不会在每次程序启动时都重新编译模型那样启动时间会特别长。正确做法是把编译好的engine文件保存下来部署时直接加载反序列化。加载代码很简单def load_engine(engine_file_path): with open(engine_file_path, rb) as f: serialized_engine f.read() runtime trt.Runtime(TRT_LOGGER) engine runtime.deserialize_cuda_engine(serialized_engine) return engine重点在于理解engine文件是和特定GPU绑定的。你在A卡上编译出来的engine拿到B卡上反序列化很可能要么报错要么性能不是最优。原因是tactic选择和GPU架构强相关。正确流程是拿到目标部署机或者同规格GPU的机器上重新编译一遍。我之前吃过亏把一台测试机上编译好的engine拷贝到生产服务器结果直接抛错后来老老实实在生产机上建了个编译任务。4.4 批处理一个stream跑起来Engine编译完成后推理阶段的性能除了受算子和精度影响还极大依赖批处理大小和CUDA Stream的使用方式。很多入门教程只演示了batch1的推理这没错但生产中如果要追求吞吐一定要充分利用batch维度。TensorRT的Engine在编译时已经考虑了batch大小的优化特别是动态shape场景下batch8或16的推理吞吐通常远高于batch1的8倍。另外推理时一般用独立CUDA stream执行避免和默认stream冲突。代码层面需要import pycuda.driver as cuda import cuda这一步需要装上合适的CUDA Python绑定比如pycuda或cuda-python。不过这里我不展开推理细节否则内容就偏离“模型编译”这个主题了。你只需要知道编译阶段决定性能上限推理阶段的任务是如何把上限发挥出来。5. 实测中常见的问题与排查技巧实录5.1 最常见的两类报错我在帮很多朋友排查第一个trt程序时发现报错高度集中在两类。第一类ONNX解析失败。特征是最开始那几行日志一片红然后告诉你某个算子找不到或者opset版本太新/太旧。解决办法是先看报错里点名的是哪个算子去查这个算子是不是有第三方实现如果只是opset版本问题用onnx-simplifier做一次简化重构或者从模型导出端调整opset。这里尤其推荐用onnxsim处理一遍再进TensorRT能省掉很多不必要的麻烦。第二类显存不足。这类问题多见于大模型和动态shape。如果报错信息里有out of memory字样大概率是workspace设得太大或者profile太多。我的习惯是先把workspace降到512MB左右再逐步加大直到找到一个平衡点。另外如果是动态shape场景检查一下profile.max_shape是不是设得比实际需要大太多了max_shape直接决定了预留显存上限。5.2 怎么判断编译是否成功这里给一张速查表方便你往后排查现象可能原因解决方案parser解析时报大量错误ONNX算子不兼容或图结构太复杂用onnxsim简化或检查导出时的opset编译过程变慢到几十分钟workspace设置过大或模型很大降低workspace或拆分模型开了FP16后精度崩了某些层不适合FP16设置精度范围强制部分层走FP32Engine在其他机器上加载失败GPU架构不匹配在目标卡或同架构卡上重新编译编译一直OOMworkspace和profile过大减小workspace、精简约max shape编译一分钟内完成推理也没变快可能没找到最优tactic增大workspace重新编译5.3 提升编译成功率的小技巧如果模型结构比较复杂我强烈建议在进TensorRT之前先用onnx-simplifier把ONNX模型做一次图简化。很多时候模型导出时带了大量无用节点比如一堆Identity、Cast、Transpose这些节点本身不耗时但会让parser解析变慢也可能引发奇怪错误。简化之后模型结构干净不少TensorRT的层融合也能做得更彻底。另外编译环境最好开一个独立终端不要同时跑多个占用显存的进程。Builder做tactic扫描时需要真实验证候选kernel显存占用在编译高峰可能超过预期。我遇到过两次因为同时开着可视化工具喂显存结果编译中途失败的情况。6. 长辈式经验编译这一步值得多花时间我自己用TensorRT做了好几个模型部署项目回头看模型编译阶段其实是整个部署链路里最值得花时间的环节。很多在推理阶段怎么调都调不动的问题根子都在编译阶段没配好参数、没选对精度、没设好workspace。一个很实际的经验是先写一个一键编译脚本把参数全部显式暴露出来比如half、workspace上限、动态shape范围全都做成命令行入口。这样每次换模型、换卡、换精度只需要改参数重跑不用改代码。我后来所有项目都沿用这套编译脚本省了太多时间。还有一点编译log要养成保存的习惯。比如build_engine(model.onnx, model.engine, halfTrue)执行完除了打印结果我自己会把TensorRT日志重定向到文件里。后面如果推理出现奇怪的精度问题翻编译日志往往能定位到是哪一层被融合、哪一层走了低精度。这个习惯帮过我很多次。走到这一步你已经知道怎么写第一个trt程序也明白模型编译的过程中Builder、Network、Config、Parser各自是什么角色了。接下来最值得做的一件事就是拿一个你熟悉的小模型跑通FP32再切FP16对比两次编译出来的engine文件大小和后续推理耗时亲手感受一下“编译”这两个字的含义。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑