ONNX导出、TensorRT优化、服务化——这三件事，每个接触过Transformer落地的朋友应该都不陌生，但真正从训练好的PyTorch模型一路走到线上接口，中间的路况远比想象中复杂。我最近刚把一个基于Transformer的编码器模型完整走了一遍从导出到上线的流程，趁着印象还热，把这次实战中的方案选型、参数细节和踩过的坑一次说清楚。

这套链路的核心思路其实很朴素：训练框架追求灵活性，生产环境追求极致效率，两者之间需要一座桥。ONNX是这座桥的通用格式，TensorRT把桥上跑的车换成性能跑车，最后用服务化框架把跑车包装成随时可调用的公共交通工具。文章既适合刚接触部署的新手，也适合已经跑通过ONNX但想进一步压榨性能的朋友。

1. 部署链路整体设计与方案选型

1.1 直接拿PyTorch部署，到底差在哪

很多人一开始都会问：我直接用PyTorch的torch.jit.script或者torch.jit.trace导出模型，配合FastAPI拉个服务，不是也能跑吗？能跑，但离“生产可用”有距离。

第一个问题是动态图带来的推理开销。PyTorch在每次forward时都要重新构建计算图，虽然底层有C++运行时兜底，但相比静态图或专用推理引擎，解释和执行的开销是实打实的。第二个问题是算子级融合。GPU推理时，kernel launch和显存读写往往比计算本身更耗时。PyTorch执行一个个小算子时，每个算子都可能触发独立的kernel启动，而像TensorRT这样的引擎会把卷积、矩阵乘、激活、归一化等相邻算子融合成一个kernel，减少启动次数和中间显存读写。Transformer模型里有大量密集的矩阵乘和LayerNorm，融合收益非常明显。

还有显存占用。PyTorch的显存分配策略偏保守，训练时预留了大量缓存，推理时往往不会自动释放。生产环境里我们经常是多个模型复卡部署，显存被某个模型占了又没真正用满，很浪费。专门的推理引擎在显存管理上更精细，通常按需分配。

1.2 部署链路的整体路线图

我的整体路线分三步，每个阶段都有明确的输入输出和验收标准：

• 第一步，把PyTorch模型导出为标准ONNX文件。这一阶段的目标不是性能，而是格式兼容性和正确性。
• 第二步，把ONNX转换为推理引擎的专有Engine文件。这一阶段会做图优化、层融合、精度校准，目标是把推理速度推到极致。
• 第三步，把Engine接入统一推理服务框架，暴露HTTP/gRPC接口，支持动态Batch、并发请求、模型热加载，目标是从“能推理”变成“能服务”。

有人会跳过ONNX，直接用PyTorch导出到TensorRT，但这样会失去框架中立性。ONNX的好处是它成了一个可审计、可复用的中间产物，后续想换推理引擎、做跨团队交接、做模型可视化，都有据可依。

2. ONNX导出实战

2.1 导出前的模型结构预处理

很多人栽在第一步：拿训练时的完整模型直接export，出来的ONNX里塞了一堆训练才需要的分支，比如dropout、梯度计算相关的算子，还有loss计算。这些不仅增加体积，还会导致TensorRT转换失败。

正确做法是导出前先剥掉推理不需要的部分。以典型的Transformer编码器模型为例，训练时输入是input_ids, attention_mask, token_type_ids，输出是logits和hidden_states。推理时我们通常只需要logits，hidden_states可以去掉。更关键的是要把dropout层设为eval模式，或者在forward里直接跳过dropout，否则导出的ONNX会包含随机的RNG算子，每次推理结果都会有随机差异。

我给自己的代码设了一个规矩：单独写一个inference_model类，只保留embedding、Transformer编码层、最后的分类头，forward里只返回logits。宁可多写几行，也不要为了省事直接export原模型。下面是一个标准的导出脚本框架：

import torch class TransformerInference(torch.nn.Module): def __init__(self, model): super().__init__() self.backbone = model self.dropout = torch.nn.Identity() # 推理时直接跳过 def forward(self, input_ids, attention_mask): last_hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask) logits = last_hidden[0] # 自行加分类头或者从原模型取 return logits model = torch.load("checkpoint.pt", map_location="cpu") model.eval() infer_model = TransformerInference(model).eval() dummy_input = { "input_ids": torch.randint(0, 20000, (1, 128), dtype=torch.long), "attention_mask": torch.ones(1, 128, dtype=torch.long), } torch.onnx.export( infer_model, (dummy_input["input_ids"], dummy_input["attention_mask"]), "transformer.onnx", opset_version=13, do_constant_folding=True, input_names=["input_ids", "attention_mask"], output_names=["logits"], dynamic_axes={ "input_ids": {0: "batch_size", 1: "sequence_length"}, "attention_mask": {0: "batch_size", 1: "sequence_length"}, "logits": {0: "batch_size", 1: "sequence_length"}, }, ) print("export finished")

2.2 动态维度到底该不该开

dynamic_axes是导出时最容易纠结的参数。固定维度导出最简单，转换TensorRT也最省事，性能最好。但线上请求的sequence length是变化的，固定维度意味着要么padding到固定长度，浪费算力，要么遇到超长文本直接报错。

我的建议是：对外服务的输入长度上限必须先在业务侧定清楚。比如我们约定最大长度512，那么动态轴只需要把batch_size设为动态，sequence_length可以固定为512，这样能把padding浪费控制在一个可接受范围，同时保留batch合并能力。如果确实需要变长序列，ONNX导出时两个维度都标为动态，但后面在TensorRT里需要做优化，需要额外配置profile。

动态维度不是免费的午餐，它会降低引擎的编译优化空间。TensorRT需要针对不同shape创建profile，动态范围越大，可选优化越折中。所以选型时优先考虑业务维度，而不是一味追求“全动态”。

2.3 ONNX Runtime本地验证，这一步不能跳过

导出后我习惯先跑一遍ONNX Runtime的CPU或GPU推理，确认数值和PyTorch输出基本一致，再继续往后做TensorRT。这一步能过滤掉大部分算子兼容性问题。

import onnxruntime as ort import numpy as np sess_options = ort.SessionOptions() sess = ort.InferenceSession("transformer.onnx", sess_options=sess_options, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) input_ids = np.random.randint(0, 20000, (2, 128)).astype(np.int64) attention_mask = np.ones((2, 128), dtype=np.int64) outs = sess.run(["logits"], {"input_ids": input_ids, "attention_mask": attention_mask}) # 与PyTorch结果对比 torch_input_ids = torch.from_numpy(input_ids) torch_attention = torch.from_numpy(attention_mask) with torch.no_grad(): ref = infer_model(torch_input_ids, torch_attention).numpy() print("max abs diff:", np.abs(outs[0] - ref).max())

只要最大绝对误差在1e-3量级以内，基本可以认为导出是正确的。如果误差很大，别急着调TensorRT，先检查模型结构、eval模式、输入dtype、dynamic_axes是否匹配。

2.4 ONNX导出阶段经典问题速查表

这些坑几乎每个项目都会碰到一两个，尤其是“导出成功但推理结果不一致”这种，排查思路要从模型的随机性入手：确认模型处于eval模式，确认没有dropout，确认输入顺序一致。

3. TensorRT优化与精度校准

3.1 TensorRT里几个绕不开的概念

第一次用TensorRT的人会被builder、network、profile、engine、context这些概念绕晕。我用自己的话捋一遍逻辑：

• Builder是全局构造接口，负责配置参数（精度、工作区大小、插件支持等）。
• Network是解析后的计算图，可以看成TensorRT内部版的ONNX。
• Profile用于描述动态shape的取值范围，每个profile包含min/opt/max三组shape。
• Engine就是最终的优化产物，它是序列化后的二进制文件，包含了kernel代码和权重。
• Context是engine的执行实例，每次推理前需要为输入输出分配合适的显存buffer。

实际开发中，大多数人不会直接调Network API去搭图，而是让TensorRT去解析ONNX，然后自动优化。所以我们工作的重心其实是：把ONNX喂进去，配置好Build参数，然后重复试错直到精度和性能都达标。

3.2 用trtexec快速验证可转换性

在写Python API之前，先用trtexec做个快速验证是非常聪明的做法。它能在不写代码的情况下测试ONNX能不能转成engine、默认性能如何。

trtexec --onnx=transformer.onnx \ --saveEngine=transformer_fp16.engine \ --fp16 \ --workspace=4096 \ --minShapes=input_ids:1x128,attention_mask:1x128 \ --optShapes=input_ids:4x128,attention_mask:4x128 \ --maxShapes=input_ids:8x128,attention_mask:8x128

如果这个命令能生成engine，说明ONNX文件在基础层面没问题。如果报错，就去看log里的unsupported op或shape推论失败信息。通常大部分问题都能在ONNX层面通过算子替换解决，不需要去写TensorRT插件。

3.3 Python API构建Engine的核心流程

生产环境里我更喜欢用Python API，因为可以更方便地集成进现有的部署脚本。一个典型的构建流程如下：

import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(TRT_LOGGER) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) profile = builder.create_optimization_profile() parser = trt.OnnxParser(network, TRT_LOGGER) with open("transformer.onnx", "rb") as f: parser.parse(f.read()) config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * 1 << 30) # 4GB config.builder_optimization_level = 3 profile.set_shape("input_ids", (1, 128), (4, 128), (8, 128)) profile.set_shape("attention_mask", (1, 128), (4, 128), (8, 128)) config.add_optimization_profile(profile) if builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) engine = builder.build_serialized_network(network, config) with open("transformer.engine", "wb") as f: f.write(engine)

这里的EXPLICIT_BATCH必须打开，因为我们的ONNX导出的batch是动态轴，不能用隐式batch模式。set_memory_pool_limit之前叫set_workspace_size，不同版本API有差异，建议按实际版本文档来。这个参数控制的是构建时的临时工作区上限，不是推理显存，别开太大导致显存溢出。

3.4 FP16和INT8量化怎么选

Transformer模型部署，FP16几乎是必选项。模型权重和激活通常是FP32占大头，转成FP16后显存直接减半，计算吞吐也接近翻倍。精度损失方面，对于分类、排序这类的任务，FP16带来的误差通常不影响最终结果。但有一个容易踩的坑：LayerNorm中的均值方差计算对精度敏感，FP16下偶尔会出现数值细微抖动。如果任务对输出差异极度敏感（比如某些回归分数），建议保留部分层在FP32，TensorRT支持per-layer精度设置，但代价是调优成本。

INT8量化收益更大，但需要仔细校准。校准的本质是收集训练集或验证集在模型各层激活值的分布，然后计算一个阈值，把FP32激活映射到INT8范围。选校准数据时一定要贴近真实业务分布。我见过有人直接拿随机噪声做校准，结果精度崩到没法用。

一个可行的INT8流程是：

config.set_flag(trt.BuilderFlag.INT8) config.set_calibrator(MyCalibrator(calibration_images, batch_size=8))

校准数据的数量不用太多，几百到一千个样本足够。但覆盖要全，各种长度、各种难度、各种典型slice都要有。校准过程本身会跑一遍前向，时间可能比构建还久，耐心等。

3.5 性能调优的几个实测心得

• 选择合适的builder_optimization_level。这个参数从0到3，级别越高，允许的图优化越激进，但构建时间也越长。生产环境固定用3，构建过程多花几分钟是值得的。
• batch和并发怎么权衡。Transformer在batch=1时，计算量相对小，瓶颈在kernel launch和显存延迟；提高batch后，计算密集度增加，整体吞吐更高。但如果同一时间只有零散请求，盲目加大batch反而增加排队延迟。所以服务化层要配合动态batch和timeout策略，不能只看engine的单batch吞吐。
• 显存回收问题。engine构建完，建议调用torch.cuda.empty_cache()或者显式释放builder和parser，否则构建期的临时显存会一直占着，影响后续服务推理。

4. 服务化部署实践

4.1 服务化方案怎么选

Engine做好之后，下一步是考虑怎么把模型暴露给上游调用。常见选择有三类：

我一般推荐直接用统一推理服务框架，尤其当你有多个模型、模型版本需要切换、流量波动比较大的场景。它的动态batching机制能直接把多个小请求拼接成一个batch，配合TensorRT的强计算能力，吞吐提升非常明显。自建API的方案只适合模型数量少、并发几百以内的demo级服务。

4.2 模型仓库目录结构设计

统一推理服务框架对模型仓库结构有严格约定。我以一个名为transformer_encoder的模型为例，目录结构如下：

model_repository/ └── transformer_encoder/ ├── 1/ │ └── model.plan ├── config.pbtxt └── labels.json

1是版本号，目录只放TensorRT生成的engine文件，命名固定为model.plan。config.pbtxt用来描述模型输入输出、格式和后端配置。

4.3 一份可直接用的config.pbtxt

这份配置我踩过不少坑，最重要的是把input和output的dims、format和name跟导出时的ONNX张量完全对齐。如果ONNX输入名是input_ids，这里就必须是input_ids，大小写都别改。

name: "transformer_encoder" backend: "tensorrt" max_batch_size: 0 # 我们用动态shape，不启用静态batch input [ { name: "input_ids" data_type: TYPE_INT64 dims: [-1, 128] }, { name: "attention_mask" data_type: TYPE_INT64 dims: [-1, 128] } ] output [ { name: "logits" data_type: TYPE_FLOAT dims: [-1, 512] } ] instance_group [ { count: 1 kind: KIND_GPU gpus: [0] } ] dynamic_batching { preferred_batch_size: [1, 4, 8] max_queue_delay_microseconds: 100 }

这里的max_batch_size:0和动态shape配合，推理服务框架会直接把请求发给TensorRT后端，由profile决定shape。max_queue_delay_microseconds控制请求在队列里最多等多久，再凑包发送给模型。这个值不能设太大，否则小流量场景下延迟会莫名其妙变高。我通常设100微秒到500微秒，具体要压测后调整。

4.4 编写客户端调用与性能压测

客户端调用以gRPC为例，比HTTP性能好。完整客户端代码略长，核心逻辑是构造输入字符串、填request、解析response。我平时先用一个简单的Python脚本做正确性和延迟验证：

import tritonclient.grpc as grpcclient from tritonclient.utils import np_to_triton_type import numpy as np client = grpcclient.InferenceServerClient(url="localhost:8001") inputs = [ grpcclient.InferInput("input_ids", [1, 128], "INT64"), grpcclient.InferInput("attention_mask", [1, 128], "INT64"), ] inputs[0].set_data_from_numpy(np.random.randint(0, 20000, (1, 128)).astype(np.int64)) inputs[1].set_data_from_numpy(np.ones((1, 128), dtype=np.int64)) result = client.infer("transformer_encoder", inputs) logits = result.as_numpy("logits") print(logits.shape)

压测时不要只测P99延迟，要同时关注吞吐。我惯用的方法是固定并发数，先跑预热，等显存和缓存稳定后记录三个指标：平均延迟、P95延迟、吞吐量(requests/sec)。调优时要看延迟和吞吐的平衡，单纯追求最低延迟可能导致吞吐很低，生产上反而容易把上游打挂。

4.5 服务化层面的高可用细节

• 多模型实例：如果单GPU显存够，可以在config.pbtxt里设置count:2，让框架开两个engine实例，处理并发请求能力直接翻倍。
• 优雅退出：推理服务框架需要处理SIGTERM信号，保证正在处理中的请求不中断，新请求不进入。
• 模型热更新：利用模型仓库的版本管理，先把新版本engine放到2/目录，等预加载完成后再切流量，回滚只需删版本目录。

5. 常见问题与避坑实录

5.1 ONNX转TensorRT时最容易出现的三种报错

第一种，张量shape推导失败。常见于自定义op或者某些动态shape场景。建议在ONNX里先固定中间层shape，转成功后再逐步放宽动态约束。第二种，不支持的数据类型。PyTorch里torch.long对应ONNX的INT64，TensorRT有些层不支持INT64，比如某些softmax或layer_norm捷径。解决办法是在导出时把冗余的long类型转成int32，或者干脆确保输入就是int64但内部计算尽量用float。第三种，插件缺失。如果ONNX里有自定义算子，TensorRT无法识别，会报“Plugin not registered”。先看能否用等价基础算子替换，实在不行再写自定义插件。

5.2 精度不一致问题的一次完整排查过程

我碰到过一次FP16 engine的输出和ONNX Runtime偏差比较大，不是个小尾巴误差，而是分类label都变了。一开始怀疑是量化问题，后来一步步排查发现是LayerNorm在FP16下分母计算溢出。解决办法有两种：一是把LayerNorm单独保留为FP32，二是对LayerNorm的输入做clip，让方差不会过小。最后我选了后者，改动小，性能几乎没有损失。

5.3 服务化部署后的监控三件套

部署上线还没完，监控跟不上等于裸奔。我的基础监控有三个维度：

• 推理服务本身的指标：请求延迟、吞吐、动态batch调度情况、GPU利用率。
• 资源指标：显存使用量、GPU温度、显存带宽。Transformer的显存占用相对稳定，如果显存异常增长，多半是内存泄漏或engine构建残留。
• 业务质量指标：输出logits分布的漂移、线上真实输入的有效长度分布。这两个指标能帮你判断是否需要对动态shape的profile范围进行调整。

有一次我们线上忽然爆出一批P99延迟陡增，后来调监控发现是输入sequence length平均长度变大，导致TensorRT走到了opt和max之间的非最优优化区间。解决方式就是把profile的opt由4x128调整到8x224，延迟立刻回落。

5.4 给后来者的一些小建议

如果让我重新做一遍，会在最开始就明确三个边界：最大输入长度、最大batch、可容忍的延迟上限。这三个数值决定了ONNX导出时的dynamic_axes怎么配，决定了TensorRT的profile怎么配，也决定了服务化能不能用动态batch。哪里有犹豫，就先拿线上业务历史数据画个分布图，用数据说话比拍脑袋强。

另外，所有脚本都要带上可复现性设计，ONNX导出时记录PyTorch版本、opset版本；TensorRT构建时记录builder版本、engine构建参数；服务化时记录模型版本和配置文件。我吃过没记录的亏，出问题后连“这个engine是用哪一版ONNX构建的”都不知道，排查效率极低。现在我会把所有构建信息写进一个build_info.json，跟engine一起归档。

最后再分享一个实战里很讨巧的习惯：每次ONNX导出后，先跑一遍onnx.checker.check_model，再用onnx.shape_inference.infer_shapes看中间张量shape，能提前发现很多人眼看不出来的维度广播问题。部署链路越长，前期的“笨功夫”越值钱。