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。推理时我们通常只需要logitshidden_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_idsinput_ids, attention_maskattention_mask) logits last_hidden[0] # 自行加分类头或者从原模型取 return logits model torch.load(checkpoint.pt, map_locationcpu) model.eval() infer_model TransformerInference(model).eval() dummy_input { input_ids: torch.randint(0, 20000, (1, 128), dtypetorch.long), attention_mask: torch.ones(1, 128, dtypetorch.long), } torch.onnx.export( infer_model, (dummy_input[input_ids], dummy_input[attention_mask]), transformer.onnx, opset_version13, do_constant_foldingTrue, 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_optionssess_options, providers[CUDAExecutionProvider, CPUExecutionProvider]) input_ids np.random.randint(0, 20000, (2, 128)).astype(np.int64) attention_mask np.ones((2, 128), dtypenp.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导出阶段经典问题速查表现象可能原因解决思路导出时报Unsupported operatorPyTorch版本与opset版本不匹配换个opset版本或升级PyTorchONNX Runtime推理结果全零或NaN模型里有自定义算子或动态shape处理不当用torch.onnx.export的custom_opsets或改模型结构输出中包含了hidden_states但根本没用到forward返回了多个值裁剪输出只保留最终logits模型体积过大权重没有冻结或包含了优化器状态加载state_dict后新建模型不要直接导出checkpoint动态维度导出了但Runtime推理报shape不一致dynamic_axes写错导致某个中间Tensor被固定用onnxruntime的shape inference检查中间节点shape这些坑几乎每个项目都会碰到一两个尤其是“导出成功但推理结果不一致”这种排查思路要从模型的随机性入手确认模型处于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 --onnxtransformer.onnx \ --saveEnginetransformer_fp16.engine \ --fp16 \ --workspace4096 \ --minShapesinput_ids:1x128,attention_mask:1x128 \ --optShapesinput_ids:4x128,attention_mask:4x128 \ --maxShapesinput_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下偶尔会出现数值细微抖动。如果任务对输出差异极度敏感比如某些回归分数建议保留部分层在FP32TensorRT支持per-layer精度设置但代价是调优成本。INT8量化收益更大但需要仔细校准。校准的本质是收集训练集或验证集在模型各层激活值的分布然后计算一个阈值把FP32激活映射到INT8范围。选校准数据时一定要贴近真实业务分布。我见过有人直接拿随机噪声做校准结果精度崩到没法用。一个可行的INT8流程是config.set_flag(trt.BuilderFlag.INT8) config.set_calibrator(MyCalibrator(calibration_images, batch_size8))校准数据的数量不用太多几百到一千个样本足够。但覆盖要全各种长度、各种难度、各种典型slice都要有。校准过程本身会跑一遍前向时间可能比构建还久耐心等。3.5 性能调优的几个实测心得选择合适的builder_optimization_level。这个参数从0到3级别越高允许的图优化越激进但构建时间也越长。生产环境固定用3构建过程多花几分钟是值得的。batch和并发怎么权衡。Transformer在batch1时计算量相对小瓶颈在kernel launch和显存延迟提高batch后计算密集度增加整体吞吐更高。但如果同一时间只有零散请求盲目加大batch反而增加排队延迟。所以服务化层要配合动态batch和timeout策略不能只看engine的单batch吞吐。显存回收问题。engine构建完建议调用torch.cuda.empty_cache()或者显式释放builder和parser否则构建期的临时显存会一直占着影响后续服务推理。4. 服务化部署实践4.1 服务化方案怎么选Engine做好之后下一步是考虑怎么把模型暴露给上游调用。常见选择有三类方案优势劣势适用场景自建HTTP服务FastAPI等灵活可控依赖少动态batch、并发管理要自己写小流量、快速验证ONNX Runtime Serving直接用ONNX模型部署简单优化能力弱不支持TensorRT engine对性能要求不高统一用ONNX统一推理服务框架如Triton支持多framework后端、动态batching、模型热加载、并发调度成熟部署运维有学习成本生产环境首选我一般推荐直接用统一推理服务框架尤其当你有多个模型、模型版本需要切换、流量波动比较大的场景。它的动态batching机制能直接把多个小请求拼接成一个batch配合TensorRT的强计算能力吞吐提升非常明显。自建API的方案只适合模型数量少、并发几百以内的demo级服务。4.2 模型仓库目录结构设计统一推理服务框架对模型仓库结构有严格约定。我以一个名为transformer_encoder的模型为例目录结构如下model_repository/ └── transformer_encoder/ ├── 1/ │ └── model.plan ├── config.pbtxt └── labels.json1是版本号目录只放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(urllocalhost: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), dtypenp.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的INT64TensorRT有些层不支持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能提前发现很多人眼看不出来的维度广播问题。部署链路越长前期的“笨功夫”越值钱。