最近在捣鼓一个部署需求训练好的 PyTorch 模型要接进服务端做实时推理模型本身不算复杂但麻烦在它同时要吃两路输入其中一路还是变长序列。本来以为 PyTorch 模型导出 ONNX 格式就是一行torch.onnx.export的事真正上手才发现多参数输入怎么传、动态维度怎么标、输出怎么验每一步都有不少细节文档里的示例又大多是单输入、全静态 shape 的“玩具代码”拿来套真实模型经常翻车。折腾完这一圈我决定把 PyTorch 模型导出 ONNX 格式、支持多参数与动态输入这件事完整记录下来。这篇文章不会只给一个能跑的最小示例而是会从部署需求倒推导出方案把多输入、多输出、动态轴的配置逻辑、部署端验证方法、以及我实际踩过的几个高频坑全部拆开来讲适合那些已经会用 PyTorch 训练模型、正准备往推理端迁移的开发者参考。1. 为什么部署场景绕不开 ONNX多参数和动态输入到底难在哪1.1 ONNX 在部署链路里的位置先明确一个概念PyTorch 模型默认是“带着 Python 运行时”的权重、前向逻辑、算子实现全部和 Python 环境绑在一起。你训练完的模型想放到服务端总不能要求对方环境里也装一套 PyTorch、还要 import 你的模型文件。更现实的问题是纯 PyTorch 做 CPU/GPU 推理性能往往不是最优很多优化手段在 Python 动态图模式下根本施展不开。ONNX 的价值就在这里它是一种开放、中立的计算图表示格式把 PyTorch 模型里的张量流动、算子调用全部固化成一张静态图。有了这张图部署端可以交给 ONNX Runtime、TensorRT 这类专门优化过的推理引擎去加载执行不依赖 Python也不用 import 原来的训练代码。但 ONNX 图本质是“静态”的建图那一刻输入输出形状、算子组合、图结构基本都定下来了。PyTorch 模型里有很多“动态”的东西——Python 的 if、for 循环、字典传参、基于张量 shape 动态构造别的张量——这些在导出时要么被固化成常量要么需要你手动改成 ONNX 能理解的算子。多参数和动态输入的麻烦正是从这里冒出来的。1.2 多参数不只是“多传几个 tensor”的事单输入模型好办torch.onnx.export(model, x, ...)第二个参数塞一个张量就完事。多输入模型就麻烦一点你的forward函数可能有input_ids、attention_mask两个张量输入导出时第二个参数需要传一个 tuple(input_ids, attention_mask)还要告诉导出器每个输入叫什么名字。名字这事看着不起眼部署端跑 ONNX Runtime 的时候喂数据就是通过名字传给模型的名字对不上直接报错。这里面的隐含规则是ONNX 图的输入是一个固定列表每个输入有自己独立的名称和形状描述。PyTorch 模型 forward 里有多少个参数导出图里就有多少个输入节点顺序必须一一对应。你如果习惯了在 Python 里用关键字传参那套导出时很容易搞混顺序。而真正到了部署端输入名就是你的 API 契约一个字母写错运行时立刻翻脸。1.3 动态输入难在“有些维度不能写死”默认情况下torch.onnx.export会按照你传入样例的具体 shape 来固定模型输入输出。比如你拿一个[1, 128]的输入做导出得到的 ONNX 图里输入维度就写死成1和128。线上推理时如果来一个[8, 128]的 batch直接报维度不匹配。更隐蔽的是有些模型内部的算子会把某个动态维度“悄悄”固化掉。比如你在 forward 里写了x.shape[0]去创建另一个张量这个操作在导出时很可能被折叠成一个常量就算你外面把 batch 维标成动态中间运算还是按固定数字走最终导出的图是“假动态”。要避免这些坑就得在动手导出之前先想清楚三件事模型里有没有不可导出的 Python 控制流输入输出里哪些维度必须是动态的部署端到底需要哪种形状的输入这三件事没理清后面导出十个有八个要返工。2. 导出前的参数准备torch.onnx.export 的每一个关键设置2.1 一个最简可用的多输入导出示例先看一个最基础的双输入模型导出代码后面所有内容都从这里展开。假设我有一个简单的双输入分类模型特征向量和一个额外的长度信号import torch import torch.nn as nn class DualInputModel(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(768, 256) self.fc2 nn.Linear(256, 2) def forward(self, feature, lengths): # feature: [batch, 768] # lengths: [batch]这里只是演示多输入 x self.fc1(feature) x x lengths.unsqueeze(1) return self.fc2(x) model DualInputModel() model.eval() feature torch.randn(4, 768) lengths torch.randn(4) torch.onnx.export( model, (feature, lengths), # 第二个参数传 tuple dual_input_model.onnx, input_names[feature, lengths], # 与 forward 参数顺序对应 output_names[logits], opset_version17, dynamic_axes{ feature: {0: batch}, lengths: {0: batch}, logits: {0: batch} } )这段代码里第二个参数(feature, lengths)就是模型 forward 的完整入参。PyTorch 文档里常写的model和x只是最简单情况真实模型只要你把 forward 需要的参数按顺序打包成 tuple 就行。2.2 opset_version 选型能力与兼容性的取舍opset_version是很多人容易忽略但又很关键的参数。它决定了导出的 ONNX 图使用哪个版本的算子集合。新版 opset 支持更多高level 算子导出的图通常更简洁但部署端的 ONNX Runtime 如果版本偏旧可能不认识新算子加载阶段直接报错。我自己的经验是先确认部署端用的 ONNX Runtime 版本支持到哪个 opset再回头定导出参数。比如 ONNX Runtime 1.16 左右对 opset 17 支持得已经很好而更老的 1.10 可能对 opset 13 更稳。如果你完全不确定部署环境保守选opset_version13是一个比较中庸的选择兼容性好代价是某些复杂算子可能会被拆成多个基础算子。这里有一个常见误区opset 越高不一定越好。新版 opset 引入的算子虽然多但推理引擎的优化覆盖可能没跟上。我自己遇到过类似情况用 opset 17 导出的图在 ONNX Runtime 上能跑但用 TensorRT 转换时反而有几个算子不支持降到 opset 13 用等价子图替换后反而顺利了。2.3 input_names 和 output_names不是随便起名字输入输出命名看似只是字符串实际上影响后续所有环节。ONNX Runtime 推理时是session.run(None, {input_ids: tensor})这个字典的 key 就是input_names里定义的值。如果导出时叫input部署端也必须是input两边差一个字符都跑不起来。我习惯在命名时直接用模型 forward 形参的名字避免后期维护时对不上。输出名也不要偷懒output、output1这种名字在调试多输出模型时非常考验记忆力。我会根据业务含义命名比如分类模型叫logits检测模型叫boxes和scores。另外input_names 的列表顺序必须和导出时传入的 tuple 顺序一致。前者是给 ONNX 图的输入起名字后者是告诉导出器“这个位置对应 forward 的第几个参数”两者错位会导致图的输入语义完全混乱。3. 多参数输入与多输出导出的完整姿势3.1 多 Tensor 输入的传参方式上一章的示例已经展示了双 tensor 输入的基本写法这里我再展开讲讲更复杂的场景。假设你的模型 forward 长这样def forward(self, input_ids, attention_mask, token_type_ids):导出时就这么写torch.onnx.export( model, (input_ids, attention_mask, token_type_ids), model.onnx, input_names[input_ids, attention_mask, token_type_ids], output_names[logits], opset_version17 )注意input_ids、attention_mask、token_type_ids三个张量的顺序必须和 forward 定义一致。如果你在导出时代入的样例张量用了别的前向计算方式比如第一个参数其实是 token_type_ids那导出的模型在部署端就会张冠李戴。这里面还有一个容易踩的细节传给torch.onnx.export的样例张量应该使用和真实推理时一致的 dtype。PyTorch 的nn.Linear默认输入是 float32而 tokenizer 出来的 input_ids 是 int64。如果你偷懒用一个随机 float 张量代替 int64 输入做导出虽然有时候也能成功但图上记录的输入类型是 float部署端按 int64 喂数据就会报类型不匹配。3.2 多输出场景检测头和多任务模型多输出模型在实际部署里非常常见典型代表是目标检测模型一个 forward 同时返回类别分数和边界框坐标。PyTorch 的 forward 可以返回 tuple也可以返回一个 listONNX 导出后这两种情况都会被展开成多个输出节点。class Detector(nn.Module): def __init__(self): super().__init__() self.backbone nn.Linear(64, 128) self.cls_head nn.Linear(128, 4) self.box_head nn.Linear(128, 4) def forward(self, x): feat self.backbone(x) scores self.cls_head(feat) boxes self.box_head(feat) return scores, boxes model Detector() torch.onnx.export( model, (torch.randn(2, 64),), detector.onnx, input_names[input], output_names[scores, boxes], dynamic_axes{ input: {0: batch}, scores: {0: batch}, boxes: {0: batch} }, opset_version17 )部署端拿到的是一个列表列表顺序对应output_names的顺序也就是 forward 返回 tuple 的顺序。这一点要特别注意ONNX Runtime 的session.run返回的是 list不是 dict。你如果想按名字取输出得自己先session.get_outputs()拿到输出名列表再对应索引。多输出的动态维度不一定和输入维度一一对应。比如检测模型输出的 boxes 数量在不同图像上是可变的这时候除了标 batch 维还要给检测框数量那一维声明动态。3.3 非 Tensor 参数整数、布尔、字典怎么处理多参数输入里比多 tensor 更麻烦的是非 tensor 参数。PyTorch 模型的 forward 里经常会带一些 Python 标量或 bool 参数比如def forward(self, x, do_normalizeTrue, seq_len128):do_normalize这种 bool 在导出时会被固化为常量导出后部署端无法修改。如果你的业务逻辑确实需要在推理时切换是否 normalize就不能把它当作普通 Python 标量留在 forward 里。我的建议是把这类参数改造成 tensor 输入。比如do_normalize可以接收一个 int64 张量[0]或[1]模型内部再做判断seq_len这种长度参数直接作为 int64 tensor 输入部署端可以传任意长度。改造 forward 确实要多写几行代码但换来的是运行时灵活性值得。还有一个几乎必踩的坑forward 里使用**kwargs或 dict 传参。ONNX 导出器不认识 dict 结构你如果有一个config字典决定模型内部行为导出时很容易报错或静默固化。处理办法是重写 forward把所有需要的参数显式列出来再在导出前把字典拆开。4. 动态轴配置让 batch、序列长度、图像尺寸都随意变化4.1 dynamic_axes 的两种写法dynamic_axes是支持动态输入的入口。最简单的写法是这样的dynamic_axes{input: [0, 1]}意思是 input 的第 0 维和第 1 维都是动态的。这种写法能跑但可读性太差你根本不知道第 0 维代表的含义到底叫batch还是叫别的什么。更推荐的是带维度名的写法dynamic_axes{ input_ids: {0: batch, 1: seq_len}, attention_mask: {0: batch, 1: seq_len}, logits: {0: batch, 1: seq_len} }维度名可以是任意字符串不影响算子计算逻辑但对后续调试和跨部门协作非常有帮助。看到batch就知道这一维是批次大小看到seq_len就知道是序列长度。这里还有一个容易误解的地方如果你给两个输入的第 1 维都命名为seq_len比如 input_ids 和 attention_mask这意味着在导出器看来这两个维度的符号是同一个。它们在实际运行时必须相等——这恰好符合 NLP 任务里 mask 和 ids 长度一致的约束。反过来如果你希望两个维度可以不同就应该用不同的维度名。4.2 典型场景配置变长序列、动态 batch、多尺度图像先看 NLP 变长序列的典型配置。输入是[batch, seq_len]的 token id输出是[batch, seq_len, num_classes]的 logitstorch.onnx.export( model, (input_ids, attention_mask), ner_model.onnx, input_names[input_ids, attention_mask], output_names[logits], opset_version17, dynamic_axes{ input_ids: {0: batch, 1: seq_len}, attention_mask: {0: batch, 1: seq_len}, logits: {0: batch, 1: seq_len} } )这样导出的模型部署端可以接受[1, 128]、[8, 512]这类不同 batch 和序列长度的输入在线服务非常需要这种灵活性。再看 CV 目标检测的多尺度配置。图像输入一般是[batch, 3, height, width]输出 boxes 可能是[batch, num_dets, 4]dynamic_axes{ image: {0: batch, 2: height, 3: width}, boxes: {0: batch, 1: num_dets}, scores: {0: batch, 1: num_dets} }num_dets是每个样本检测出的目标数量不同图不一样这一维必须动态否则导出后只能处理固定数量目标的图。4.3 怎么确认动态轴真的生效了导出完成后不能只看导出过程没报错就以为动态输入已经支持了。最直接的办法是加载 ONNX 图看输入形状描述。用 Python 写一个小脚本import onnx model onnx.load(ner_model.onnx) for inp in model.graph.input: print(inp.name, inp.type.tensor_type.shape)如果对应维度显示的是字符串比如batch、seq_len说明动态轴声明成功如果显示的是数字说明该维度被固化了。你还可以用onnxruntime分别传入不同 batch 的数据跑两次推理看是否都能正常输出这是最真实的验证。我自己检查过不少“假动态”案例图里维度名确实是字符串但模型中间某个算子把维度当常量用了。这种情况只靠看 shape 是看不出来的必须用不同形状的数据在运行时测试。4.4 动态输入的代价性能和约束动态输入不是免费的。静态 shape 的图在推理引擎里可以做很多预分配和算子融合优化动态 shape 下引擎必须为各种可能的形状准备通用方案显存分配、kernel 选择都会保守一些性能可能下降。TensorRT 这样的工具尤其明显静态 shape 可以针对特定尺寸做极致优化动态 shape 则要使用 profile 机制灵活性提升的同时优化空间变小。什么时候不需要动态轴如果你的线上服务只接受固定 batch比如batch 永远是 1或者说你的 batch size 固定是 8那完全可以把对应维度静态导出。有些场景甚至按固定序列长度 padding 好所有输入这种情况下动态轴反而是多余的还会带来性能损失。设计导出方案时动态轴只加到真正需要变动的维度即可不要图省事把所有维度都标成动态。5. 导出后的整链路验证从 onnx.checker 到运行时一致性5.1 用 onnx.checker 和图结构检查过滤低级错误导出完最先做的事是完整性检查。onnx.checker.check_model会检查图结构是否合法比如节点的输入输出是否连接到有效张量、算子参数是否完整。这一步能拦住一部分损坏的图但远远不够。我会再打印一遍图的输入输出和算子类型快速确认结构是否符合预期import onnx model onnx.load(model.onnx) onnx.checker.check_model(model) for node in model.graph.node: if node.op_type not in [Conv, Gemm, Relu, Add]: print(unexpected op:, node.op_type)如果你发现图里出现了大量 Cast、Shape、Gather 这类算子不用太惊慌这通常是导出器为了兼容 opset 做的等价转换。但如果出现了你完全没见过的自定义算子就要警惕部署端能不能支持。5.2 用 ONNX Runtime 跑一次真实推理结构检查过了性能先不谈正确性必须验证。我把 PyTorch 模型和 ONNX 模型喂同样的输入对比输出张量是否一致。下面是标准的验证模板import numpy as np import onnxruntime as ort import torch # PyTorch 推理 model.eval() dummy_input { input_ids: torch.randint(0, 100, (2, 16)), attention_mask: torch.ones(2, 16, dtypetorch.long) } with torch.no_grad(): torch_outputs model(**dummy_input) if isinstance(torch_outputs, tuple): torch_outputs [t.cpu().numpy() for t in torch_outputs] else: torch_outputs [torch_outputs.cpu().numpy()] # ONNX Runtime 推理 session ort.InferenceSession(ner_model.onnx, providers[CPUExecutionProvider]) onnx_outputs session.run( None, { input_ids: dummy_input[input_ids].numpy(), attention_mask: dummy_input[attention_mask].numpy() } ) for i, (torch_out, onnx_out) in enumerate(zip(torch_outputs, onnx_outputs)): print(foutput {i}: shape {torch_out.shape} vs {onnx_out.shape}) print(max abs diff:, np.abs(torch_out - onnx_out).max())这里要注意torch_outputs转 NumPy 时最好用detach().cpu().numpy()避免梯度图干扰。对比阈值方面fp32 模型一般最大绝对误差在 1e-5 量级如果你的模型内部有大量归一化或 softmax 操作误差略大也正常可以用np.allclose(torch_out, onnx_out, atol1e-4, rtol1e-4)做判断。5.3 容易忽略的 dtype 和输入名一致性验证时最容易翻车的不是精度而是输入类型不一致。PyTorch 里 tokenizer 出来的input_ids通常是torch.int64但如果你在导出时代入的样例是torch.long导出的 ONNX 输入类型是 int64部署端如果把数据转成 float32直接报类型不匹配。输入名也是一样session.get_inputs()拿到的 name 列表和你session.run时字典 key 必须完全一致。我自己经常在导出时用了inputs_ids这种笔误命名一直到部署端排查半天才发现是名字写错了所以验证脚本里最好把输入输出名也打出来看一眼for inp in session.get_inputs(): print(inp.name, inp.shape, inp.type) for out in session.get_outputs(): print(out.name, out.shape, out.type)6. 实战中绕不开的五个高频坑与排查思路6.1 动态轴被内部常量固化的“假动态”问题这是我最想强调的一个坑。外面看 dynamic_axes 标得很完整runtime 传不同 shape 进去却报错或者输出 shape 对的但内部计算出问题。常见原因就是 forward 里用了依赖具体 shape 的 Python 操作。比如有人喜欢写mask x.new_zeros(x.size(0), x.size(1))导出的过程中这个x.size(0)和x.size(1)会被直接替换成导出时的实际数值导致图中存在隐式常量。排查思路很简单把所有用到x.shape[0]这类写法的地方都找出来改成不依赖具体数值或使用 ONNX 的动态 shape 算子的写法。比如需要生成[batch, 1]的 mask可以用torch.ones_like配合 slicing尽量让导出器把操作建模成动态的 Shape 和 ConstantOfShape 节点。判断是不是这个原因可以把导出时的 batch 设为 1部署时用 batch 4 跑报错就说明有地方固化。6.2 算子不支持和等价替换导出时报Exporting the operator aten::xxx to ONNX opset version 17 is not supported这是常见疑难。字面意思很清楚PyTorch 里的某个算子找不到对应的 ONNX 算子。通常有三个方向解决。第一种是调高 opset_version某些算子在高版本 opset 才有对应实现。第二种是改模型结构找到报错算子的实现方式用 PyTorch 基础算子重写一遍。比如有些模型里的循环控制流可以改造成矩阵运算或 Transformer 结构里的 mask 运算。第三种是使用torch.onnx.is_in_onnx_export之类的工具做导出分支让模型在导出时走一条不含不支持算子的路径。def forward(self, x, trainingNone): if torch.onnx.is_in_onnx_export(): # 导出时走简化路径 return self.simple_path(x) # 正常训练/推理路径 return self.complex_path(x)这种代码在训练阶段不生效只影响导出阶段的行为算是一种很实用的逃生通道。6.3 输入名和部署端字典 key 对不上导出模型后到部署端喂数据经常遇到Invalid Feed Input Name的报错。几乎每次都是输入名对不上。排查方式很简单打印session.get_inputs()拿到实际名称再和部署代码的字典 key 比对。由于不同版本的 ONNX Runtime 在某些场景下会自动做输入名纠正或大小写不敏感匹配本地验证没问题不代表生产环境没问题。我习惯在部署端代码里加一个启动时断言检查所有必传输入名都存在宁可启动就报错不要等到推理时才被几百毫秒的“神秘失败”折磨。6.4 多输出结果类型和顺序的误解ONNX Runtime 的session.run(None, feed_dict)返回固定顺序的 list对应output_names。很多从 PyTorch 迁移过来的人会下意识以为返回的是 dict 或者 namedtuple拿到 list 后又用名字索引结果越写越混乱。我的做法是导出前就记录一份输入输出映射表部署端代码里显式做 list 到业务变量的映射outputs session.run(None, feed_dict) logits, boxes outputs[0], outputs[1]如果输出顺序调整过记得同步修改映射别依赖“好像是这个顺序”的记忆。6.5 动态输入在 GPU 场景下的性能陷阱最后一个坑属于“导出没问题、部署后才发现”的类型。动态 shape 的 ONNX 图上 GPU 推理显存分配无法提前固定推理引擎每次都要根据当前 shape 做适配吞吐上不去。尤其在做 batch 聚合优化时服务端如果同时来了一堆不同长度的请求动态 shape 模型很难做到 batch 合并。这并不是说动态输入不好而是你要清楚它解决的问题是“形状可变”代价是牺牲一部分性能优化空间。如果线上 qps 压力很大可以反过来考虑静态导出几个常见 shape 的模型比如 seq_len 固定用 128 和 512 各导一份服务端按请求长度路由到对应模型既能享受静态优化又覆盖大部分输入分布。以上这些坑我基本都亲历过每次排查到最后都会感叹ONNX 导出这件事一半功夫在导出前的结构审查和参数设计上另一半在导出后的验证和部署联调上。真正动torch.onnx.export那几行代码反而是整个链路里最轻松的部分。