做模型部署这几年,我收到过最多的转换需求就是这一类:训练阶段用的是PyTorch,线上推理想上TensorRT,最好能从PyTorch侧直接产出trt文件,不要来回导模型,如果能顺手把INT8量化做了就更理想。这个需求看上去就是"转一下"的事,实际推进的时候却经常在版本配对、导出设置、校准数据、运行绑定这些环节上翻车。这篇文章把我完整跑通的一条链路、关键代码和踩过的坑都整理出来,包括torch.onnx.export的用法、TensorRT Python API构建引擎的流程、INT8校准器的实现,以及trt文件加载运行的全过程。文章从核心原理讲到手能落地的代码,适合那些已经会用PyTorch训练、但第一次接触TensorRT部署,或者之前试过转换但被各种报错拦住的人。
1. 先搞清楚"直接转"到底有哪几条路,再动手不迟
先从需求本身说起。标题里写的"pytorch直接转tensorrt的trt文件",我理解的核心诉求是:不要为了转TensorRT而把部署链路搞复杂,最好在PyTorch内部"一条龙"出结果。但在实际工程里,"直接转"至少对应着三条不同的技术路线,它们产出的文件和后续运行方式都不一样,这个差异如果一开始没搞清楚,后面会非常痛苦。
1.1 三条技术路线的本质差异
第一条路线是Torch-TensorRT,TensorRT生态里官方维护的PyTorch扩展。它可以接收torch.nn.Module,在模块内部把能优化的子图编译成TensorRT engine,最终交出来的仍然是一个TorchScript模块,可以用torch.jit.save保存。这个方案最贴合"从PyTorch直接转"的字面意思,适合项目代码已经基于TorchScript链路的人使用。但要注意,它最终保存的不是纯粹的.trt文件,引擎被封装在TorchScript里面。部署端如果还想用TensorRT原生的C++/Python API去管理引擎,需要额外拆解或用TorchScript加载。
第二条路线是ONNX中转,这是目前社区里资料最多、也最容易被搜到的方式。流程是先用torch.onnx.export把PyTorch模型导成ONNX,再用TensorRT的ONNX parser解析并构建出.trt序列化引擎。严格说它多了一步中间格式,不算"纯直接转",但实际部署中,最终得到的.trt文件是干净彻底的TensorRT原生引擎,C++和Python都能直接反序列化加载运行,维护灵活度最高。
第三条路线是fx前端或自己手动搭网络。对于PyTorch生态的绝大多数模型,用ONNX中转已经足够。fx前端适合算子高度定制化的场景,但需要额外处理子图替换,工程成本明显偏高。
三条路线的选择,核心区别在你要的是什么:
| 路线 | 输入 | 产出物 | 部署方式 | 上手成本 |
|---|---|---|---|---|
| Torch-TensorRT | torch.nn.Module | TorchScript模块(内含engine) | TorchScript加载或提取engine | 低,集成感强 |
| ONNX中转 | .onnx文件 | .trt序列化引擎 | TensorRT原生Runtime | 中,模板成熟 |
| fx前端/自建parser | 图结构 | .trt序列化引擎 | TensorRT原生Runtime | 高,适合特殊场景 |
1.2 什么情况下"直接转"才能落地
我个人在实际项目里更常用的是ONNX中转,因为团队里有人需要拿到.trt文件去做独立的后端服务,也有人需要把同一个onnx导给别的框架做对比实验,中间产物是通用的。如果你是多模型统一部署、要维护引擎生命周期、后续可能要上动态batch或者多stream推理,那最终得到的.trt文件比TorchScript封装更好控制。
如果团队链路完全基于PyTorch,并且没有跨平台复用ONNX的需求,Torch-TensorRT会更省事。它的"直接"体现在不用自己写parser、不用管理ONNX opset版本与TensorRT兼容性的问题,compile接口接收inputs示例后直接输出可运行的模块,这对只把TensorRT当加速后端使用的人来说确实是最平滑的。
不过有一条必须先说清楚:不管选哪条路,PyTorch里的自定义层、动态控制流、以及过于复杂的Python原生操作,都不可能被"直接转"成TensorRT原生算子。TensorRT能加速的是可以静态化的图计算部分。编译失败时错误信息往往指向某个不支持的op,此时最优解是把这个op相关的子图切分出来,用TensorRT插件机制或者子图回退去解决,而不是死磕"一键转换"。这个认知越早建立,后面排查就越有方向。
2. 版本矩阵:驱动、CUDA、TensorRT、PyTorch配对是第一道坎
转换流程调通之前,最先被坑的往往是环境。PyTorch和TensorRT各自都在快速迭代,它俩之间没有"任何版本都能随意组合"的保证。网络上能找到的教程多半写于几个月甚至一年前,版本参考意义有限,真正能少踩坑的办法是先定版本,再写代码。
2.1 一个相对稳妥的版本组合
如果你的部署机器是全新的,或者可以自由装环境,我会建议先看TensorRT版本,再决定CUDA和PyTorch。因为TensorRT的CUDA兼容范围约束更严格,PyTorch反而相对宽松,运行时检测CUDA版本不满足时一般会有明确报错。
一个相对保守的组合是:驱动版本520以上、CUDA 11.8、cuDNN 8.9、TensorRT 8.6.x、PyTorch 2.0.x。这个组合的优点在于TensorRT 8.6对ONNX opset 13到17的支持都比较成熟,PyTorch 2.0导出ONNX在export模式下的行为也比较稳定,不会出现太多dynamo相关的意外。
如果想用更新的特性,比如TensorRT 9.x引入的某些量化和稀疏优化,那可以上CUDA 12.1、TensorRT 9.0/9.2、PyTorch 2.1+,但前提是GPU必须支持对应的计算能力,建议提前用nvidia-smi确认。
| 组合 | 驱动 | CUDA | cuDNN | TensorRT | PyTorch |
|---|---|---|---|---|---|
| 保守稳定 | 525+ | 11.8 | 8.9 | 8.6.1 | 2.0.x |
| 新特性向 | 535+ | 12.1 | 8.9 | 9.2 | 2.2.x |
2.2 版本不匹配的典型表现与排查顺序
版本不匹配时,报错不会直接说"版本不对",常见的表现是:导入tensorrt时提示找不到某个cudart或nvinfer动态库;或者构建引擎时跑到一半突然崩了;又或者反序列化引擎时提示校验失败。
遇到这种问题,按下面的顺序排查,比当初我盲试要快得多:
- 先看TensorRT的Release Notes里声明的CUDA/cuDNN兼容范围。
- 用
nvidia-smi查看驱动版本与可支持的最大CUDA运行时版本。 - 用
nvcc --version和Python中torch.version.cuda分别确认两套CUDA信息。 - 检查tensorrt库的安装方式是wheel包还是tar包,tar包安装时环境变量有没有配好。
另外,pycuda的版本也要对得上。多数情况下用pip install pycuda装到的版本问题不大,但如果编译cuda扩展报错,建议直接装匹配版本的预编译轮子。调试内存绑定之前,先让import tensorrt as trt在干净环境里能稳定跑通,否则后续每个报错都可能是假象。
2.3 开始转换前,模型本身要做哪些预检
环境确认后,打开PyTorch模型,先不要急着导出。我会先做三个预检:
- 把模型切到eval模式并锁住BN和Dropout。如果训练完后BN还在跑均值方差更新,导出图的统计口径就和推理不一致,这也是INT8校准结果奇怪的常见来源。
- 确认输入张量的shape和dtype。动态batch和静态batch导出的代码写法不同,dtype若不是fp32,先统一成fp32再导出。
- 用一组真实输入跑一次forward,确认输出与预期一致。这里不是为了验证精度,而是确保模型结构本身不依赖训练阶段的中间状态。
这个预检看似多余,但它能把"模型本身就没写好"和"TensorRT转坏了"两类问题彻底分开。后面无论报什么错,都知道该往哪一层追。
3. 核心链路:PyTorch导出ONNX并用TensorRT构建trt引擎
环境就绪之后,进入正文主流程。如果走ONNX中转路线,完整链路是:PyTorch导出ONNX、TensorRT OnnxParser解析、BuilderConfig配置、构建引擎、序列化保存。下面逐个拆开讲。
3.1 导出ONNX时最容易出问题的几个设置
PyTorch导出ONNX的调用并不复杂,但这几个参数必须根据实际场景设置清楚。
import torch def export_onnx(model, dummy_input, output_path, dynamic_batch=False): model.eval() dynamic_axes = None if dynamic_batch: dynamic_axes = { "input": {0: "batch_size"}, "output": {0: "batch_size"} } torch.onnx.export( model, dummy_input, output_path, opset_version=16, input_names=["input"], output_names=["output"], dynamic_axes=dynamic_axes, do_constant_folding=True, export_params=True )opset_version的选择比较关键。TensorRT 8.6对ONNX opset 13到16的兼容性普遍可靠,opset 17之后部分新增算子TensorRT可能还不认,保守一点从16开始试。如果模型里用到了较新的PyTorch算子,导出时遇到不支持的情况,再往上调版本。
dynamic_axes的作用是声明哪些维度是动态的。注意,光在导出时声明还不够,TensorRT侧也要开启dynamic shape相关的配置,否则即使onnx里带动态维,解析到TensorRT时也可能退化成固定维度。我的习惯是:第一次转换全部用静态shape跑通,确认链路没问题后再把batch维度设成动态。静态shape下的定位成本最低,可以专心排查算子兼容性问题。
导出后先用onnx.checker.check_model校验一遍,再用onnx_shape_inference.infer_shapes做一次shape推断。这两步能提前发现不少导出阶段被隐藏的问题,比如某个节点的输出shape不明确、常量折叠后产生孤立节点等。常见的报错如"ONNX Node 'Mul' has no corresponding TensorRT Op",大多会在转换阶段追到某一个不支持的算子时暴露出来,此时再回来看ONNX图,定位会更准确。
3.2 构建引擎:从Builder到Engine的完整代码
ONNX准备完毕后,TensorRT的构建代码模板很固定,核心步骤是:创建builder、显式batch标记、解析onnx、配置、构建。我直接给一份项目里简化的构建函数。
import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def build_engine(onnx_path, engine_path, fp16=False, int8=False, calib=None): builder = trt.Builder(TRT_LOGGER) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, TRT_LOGGER) with open(onnx_path, "rb") as f: ok = parser.parse(f.read()) if not ok: for i in range(parser.num_errors): print(parser.get_error(i)) raise RuntimeError("ONNX parse failed") config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1GB,TensorRT 9.x建议改用memory pool接口 if fp16: config.set_flag(trt.BuilderFlag.FP16) if int8: config.set_flag(trt.BuilderFlag.INT8) if calib is not None: config.int8_calibrator = calib engine = builder.build_engine(network, config) if engine is None: raise RuntimeError("build engine failed") with open(engine_path, "wb") as f: f.write(engine.serialize()) return engine有几个细节值得展开一下。create_network里的EXPLICIT_BATCH标记必须加,否则沿用的是隐式batch模式,后续处理动态shape和INT8校准时会遇到一堆兼容性问题。很多人第一次转换时漏了这个标记,builder直接返回无效network,后面所有设置全都串味。
config.max_workspace_size决定构建和运行时允许使用的显存上限。给得太小会明显降低构建成功率,给得太大在显存受限的部署机上又容易资源分配失败。我一般先给1GB跑通,确认没有换层问题后,再按实际显存调整。注意TensorRT 9.x开始,这个字段有了新的等价接口,建议用set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, size)。
3.3 FP16与动态Shape:为什么建议先跑通静态再谈进阶
FP16开启的代码非常简单,config.set_flag(trt.BuilderFlag.FP16)一行就够。但精度风险是全局性的,某些数值敏感层在大模型里会出现明显退化,常见的是LayerNorm、Softmax前的sum计算。如果只是要把模型转出来跑demo,FP16是性价比最高的加速选项;如果要上线,建议用一遍验证集过一下,确认输出误差在可接受范围再开启。
动态Shape在TensorRT里不是简单声明完就结束。除了ONNX导出时声明dynamic_axes,还要在TensorRT侧配置optimization profile:
profile = builder.create_optimization_profile() # min, opt, max profile.set_shape("input", (1, 3, 224, 224), (1, 3, 224, 224), (8, 3, 224, 224)) config.add_optimization_profile(profile)这里三组shape分别是min、opt、max。跑推理时,每次输入的batch size变了,都要调用context.set_binding_shape更新输入shape,输出端的内存也要动态分配。这一步是"转出来能跑"和"实际生产可用"之间的分水岭。我的建议是:第一次落地的项目先用静态shape完成整个流程,把推理延迟、吞吐量、精度表现摸清楚,再引入动态batch,不要在转换初期同时引入两个不确定因素。
4. INT8量化:决定转换价值的核心环节
"量化int8"是整个标题里技术含量最高的部分,也是最容易拖垮进度的部分。很多人以为在BuilderFlag里加一个INT8标记就行,实际上远远不止。TensorRT做INT8量化不是简单地给权重取整,它需要在数据集上做校准,统计每一层激活值的分布,从而确定合理的量化阈值。
4.1 INT8量化的原理与TensorRT的实现方式
对称量化是把浮点数值映射到[-127, 127]区间,核心是确定缩放因子scale。对权重来说可以静态统计;对激活值来说,它的分布和输入数据强相关,不能用某一层的静态最大绝对值硬套。否则遇到长尾分布时,即便最大绝对值很大,大部分数据却集中在很小的区间,直接除以大绝对值会损失大量有效精度。这就好比一个房间最高处3米,但多数人的身高集中在1.7米附近,你用3米去做标签刻度,人群中间的差异就很难区分了。
TensorRT的解决方案是校准。校准分两步:喂一批有代表性的数据做网络前向传播,过程中统计每层激活值的分布,然后按不同策略寻找最优阈值。Python接口里常用的校准器有IInt8EntropyCalibrator2、IInt8EntropyCalibrator、IInt8MinMaxCalibrator、IInt8LegacyCalibrator。其中EntropyCalibrator2是默认推荐,它基于信息熵选择阈值,对大多数CNN模型效果不错;MinMax按最小最大绝对值直接算,简单但容易受离群点影响;Legacy是历史遗留选项,新项目基本不用。
校准器核心不是继承哪个类,而是提供正确的get_batch方法返回一批张量数据。TensorRT在校准阶段会不断调用get_batch,直到返回None。校准数据必须是推理输入对应的真实样本,不能用随机噪声替代。随机噪声可以作为构建时的占位,但用它校准出来的scale会导致模型精度崩掉。
4.2 校准器类完整实现与校准数据准备
下面是用EntropyCalibrator2做校准的完整实现。DataLoader来自PyTorch侧,校准数据经过预处理后转成numpy数组,再拷进预分配的ctypes内存。
import os import ctypes import numpy as np import tensorrt as trt class EntropyCalibrator(trt.IInt8EntropyCalibrator2): def __init__(self, dataloader, input_name, cache_file, batch_size): super().__init__() self.dataloader = dataloader self.dataloader_iter = iter(dataloader) self.input_name = input_name self.cache_file = cache_file self.batch_size = batch_size self.batch_idx = 0 self.buffer = None def get_batch_size(self): return self.batch_size def get_batch(self, names): if self.batch_idx >= len(self.dataloader): return None try: images, _ = next(self.dataloader_iter) except StopIteration: return None images = images.numpy() if self.buffer is None or self.buffer.size != images.nbytes: self.buffer = ctypes.create_string_buffer(images.tobytes(), images.nbytes) else: ctypes.memmove(self.buffer, images.tobytes(), images.nbytes) self.batch_idx += 1 return [ctypes.addressof(self.buffer)] def read_calibration_cache(self): if os.path.exists(self.cache_file): with open(self.cache_file, "rb") as f: return f.read() return None def write_calibration_cache(self, cache): with open(self.cache_file, "wb") as f: f.write(cache)有几个实现细节需要注意。get_batch返回的是一个Python列表,列表里每个元素是ctypes地址。内存必须持续有效,如果每次临时创建numpy数组再传地址,返回后内存被释放,校准阶段会读到渣值。我用了一个预分配的create_string_buffer承接数据,保证整轮校准过程中地址稳定。这是我最初用临时数组跑INT8时踩出的教训,精度差到完全没法用,查了半天才发现是校准数据在传输过程中被垃圾回收了。
校准数据的读取策略也值得认真定。校准集不是越多越好,关键是覆盖真实分布。一般一张GPU、batch size为32或64的前提下,选500到2000张有代表性的样本就够了。样本要覆盖各类输入,比如不同光照、不同目标尺度的图片,而不是只选模型训练集里最顺滑的那几百张。校准集合如果太小,得到的scale会明显偏差,典型表现是INT8模型在个别目标类别的召回显著变差。
校准cache文件要认真对待。TensorRT允许把校准结果缓存下来,下次构建直接读缓存,省去重新校准的耗时。我在项目中把这个文件当成特殊的"模型资产"管理,因为一旦缓存写了错误分布的数据,后续所有构建都会继承这个坏scale。验证缓存是否靠谱最简单的做法,是删除缓存重新校准一次,对比两次INT8模型的精度差异是否在可接受范围内。
4.3 量化后精度下降的排查重点
INT8构建成功不等于转换结束。如果量化后精度明显退化,我会按下面顺序定位:
- 先确认是"能跑但精度差"还是"某一层输出爆炸"。前者大概率是校准分布问题,后者更可能是某类敏感算子。
- 用TensorRT的层精度标记,把精度敏感层单独保留FP16或FP32:
layer.precision = trt.float16,配合layer.set_output_type设定该层输出精度。 - 尝试不同校准器,
IInt8MinMaxCalibrator虽然保守,但对某些小模型反而稳定。 - 观察是不是动态shape下校准distribution和运行shape不一致。INT8校准在网络前向时输入shape是固定的,如果后面把shape改得很大,某些层的激活分布被推到校准区间的边缘,精度自然劣化。
另一个容易被忽略的点:很多PyTorch模型在导出ONNX前带的是GELU、SiLU这类激活函数,它们在INT8下的分布形态和ReLU不一样,对称量化效果普遍更差。遇到这类模型可以选择把激活层保留FP16,或者单独做非对称量化支持。TensorRT较新版本对INT8非对称量化有更好的支持,但这需要网络本身带着Q/DQ节点,导出的ONNX结构就完全不一样了,属于另一套流程。
5. 运行trt文件:加载、绑定内存与性能验证
构建引擎只是完成了前半段,验证一个trt文件能不能正常跑,才是"并运行"这个关键词的关键。很多人在这一步用简单的numpy内存操作硬凑,结果跑出来的结果风马牛不相及,最后回头怀疑是转换阶段出了问题。先理清运行阶段的数据流:CPU上的numpy数组、GPU上的显存、引擎输入、引擎输出、复制回CPU。
5.1 序列化保存与反序列化加载的正确姿势
engine.serialize()返回的是二进制字节串,保存成.trt文件时注意一定要用二进制模式写文件。加载时,直接读文件字节再交给Runtime反序列化即可,但有几个细节:
def load_engine(engine_path): import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) runtime = trt.Runtime(TRT_LOGGER) with open(engine_path, "rb") as f: engine_bytes = f.read() engine = runtime.deserialize_cuda_engine(engine_bytes) if engine is None: raise RuntimeError("failed to deserialize engine") return engine反序列化成功只代表二进制结构完整,不代表引擎能在当前设备上运行。TensorRT引擎会在serialize时记录构建时的设备和依赖库信息,如果拿A机器构建的trt文件拿到B机器跑,经常会在deserialize阶段报错或运行时报错。条件允许时,让构建和运行的设备保持一致,或者至少保证CUDA和TensorRT版本一致。
有一种常见误解是trt文件可以像ONNX一样直接跨平台移植。实际不行。它本身就是强绑定硬件和环境的序列化产物,加速效果也是绑定在具体GPU型号上的。要在新机型部署,最靠谱的做法是在目标机器上执行一遍构建流程。
5.2 执行上下文与内存绑定的细节
拿到engine之后要创建ExecutionContext,真正执行推理靠的是context,不是engine。这一步很多人容易混。
执行时先要搞清楚输入输出bindings的索引。用engine.get_binding_index(name)获取,再查看name对应的维度信息。
context = engine.create_execution_context() input_idx = engine.get_binding_index("input") output_idx = engine.get_binding_index("output") # 静态shape场景下,可以直接用engine本身的信息分配内存 input_shape = engine.get_binding_shape(input_idx) output_shape = engine.get_binding_shape(output_idx)如果是动态shape,需要在每次推理前调用context.set_binding_shape(input_idx, (batch, channels, height, width))。设置完成后,输出的shape也要重新查询。确认shape之后,在GPU上为输入输出各分配一片显存即可。
显存分配和拷贝这一步,可以用pycuda,也可以用cupy,看团队已有的包。一个pycuda版本的执行片段:
import pycuda.driver as cuda input_size = trt.volume(input_shape) * np.dtype(np.float32).itemsize output_size = trt.volume(output_shape) * np.dtype(np.float32).itemsize input_buffer = cuda.mem_alloc(input_size) output_buffer = cuda.mem_alloc(output_size) # 每次推理前把numpy数据拷到显存 cuda.memcpy_htod(input_buffer, input_array.reshape(-1)) # 执行推理 context.execute_v2(bindings=[int(input_buffer), int(output_buffer)]) # 把结果拷回 cuda.memcpy_dtoh(output_array, output_buffer)注意execute_v2本身是同步调用。如果要用异步版本execute_async_v2,需要额外维护一个CUDA流,并在调用前确保数据拷贝和流事件同步到位。对很多项目来说,同步版本把流程跑通、把指标测明白就够了。
5.3 性能验证:该怎么对比才可信
性能验证不是简单跑一两次看时间。TensorRT引擎有warmup效应,第一次推理要加载算子kernel和建立上下文,测出来的延迟会偏高。规范的做法是:连续跑一批推理做预热,预热次数建议50到100次,然后用后面100到1000次的结果计算平均延迟和P95延迟。
如果你的目标是横向对比PyTorch GPU推理和TensorRT推理,务必要保证对比口径一致:都是单batch单流、都是相同输入数据、都不包括数据加载时间。拿PyTorch的完整推理耗时去和只统计execute_v2耗时的TensorRT对比,结论没有任何参考价值。我见过有组员汇报TensorRT比PyTorch快20倍,细问之下发现PyTorch那边统计的是包括数据预处理和CPU传输栈的全链路。
还有一点容易被忽略:INT8模型在GPU上不一定每一项都快于FP16。某些显存带宽受限的小模型,INT8的数值计算速度优势不足以抵消某些层增加的overhead,整体延迟可能只比FP16快10%甚至持平。此时不要硬上INT8,先看具体资源约束再决定。
6. 部署中实际踩过的坑:报错排查与经验沉淀
这个部分我想按排查链路写几个真正发生过的场景,只保留通用错误模式,方便你复现排查思路。
6.1 高频报错的完整排查链路
案例一:ONNX parse时提示"some tensors are not connected"。
这不是TensorRT解析器的问题,而是ONNX图本身存在未连接的张量。通常发生在PyTorch导出的子图里包含死分支或未被使用的输出。排查时先运行onnx.checker.check_model,然后在netron里看孤立节点;如果确定某个输出完全没用,就在export时去掉output_names,或者对无用输出做降维后再导出。
案例二:INT8校准阶段报"calibration data is missing"。
原因一般是config.int8_calibrator没有被赋值。注意在TensorRT Python接口中,先set_flag(trt.BuilderFlag.INT8),再把calibrator对象赋给config.int8_calibrator。如果构建流程中函数作用域没保持calibrator实例存活,也会在build阶段报缺失。Python侧要保证calibrator被引用到build_engine返回,否则垃圾回收会把它清算掉。我有一个印象很深的排错经历,就是把这行赋值写在函数里,calibrator在函数返回时被回收,构建一直报找不到校准数据。
案例三:运行阶段报"engine requires output tensor"。
常见于执行推理前没有按动态shape设置好所有绑定的shape,或者output内存没有按最新shape重新分配。按5.2节的流程,先查询输出shape再分配内存,问题就能消除。
案例四:FP16模型推理结果出现NaN。
先排查是否在PyTorch导出时用了和训练时不同的数值范围,比如输入没有归一化。如果输入范围没问题,再看LayerNorm或类似敏感层。在TensorRT中可以通过Network定义替换精度,把敏感层改成FP32输出,再重新构建引擎。这个排查过程要保留每个版本的引擎文件和对应的精度指标,否则改到后面连自己都不清楚哪一版是好的。
6.2 会写进团队文档的几个经验
第一个经验:构建流程要做成可复现的脚本,不要每次手动操作。构建引擎本身可能耗时几分钟到几十分钟,手动改一次配置重新构建一次非常低效。把onnx路径、engine路径、FP16/INT8开关、校准器、batch size都参数化,命令行一键构建。
第二个经验:trt文件要在CI流程里做健康检查。反序列化后跑一遍固定的test input,输出结果和onnx转出的参考结果做误差比对。如果误差超过预设阈值就告警。否则模型一更新,量化退化可能在几天后才会被线上监控发现,那时候定位成本就高了。
第三个经验:校准数据、校准缓存、trt文件、onnx文件要成组管理。用带哈希的目录名,模型每更新一次就生成一组新文件,不要覆盖旧的。这样线上版本回退的时候,可以精确知道当时用的是哪一份校准数据、哪一份引擎文件。这个细节在排查"为什么线上比测试慢20%"时帮过我不少,那次就是回退后发现线上跑的是旧引擎,捆绑的 workspace 配置和优化 profile 都不一样。
最后一个小建议
如果你也是第一次跑通这个链路,建议把FP32静态batch、FP16、动态shape、INT8这四个阶段当成四道关卡,每一关都保留完整的模型文件和性能记录。别想着一步到位直接上INT8动态shape,出了问题你根本说不清是校准集选坏了、profile设错了,还是某个算子精度崩了。我每次遇到转换问题,第一反应永远是"把优化开关退回到上一档,看看问题是否复现",这一步操作能省掉至少一半的无效排查时间。这套链路跑熟之后,从PyTorch到trt文件再到INT8运行,也就是一顿饭的工夫。