PyTorch多参数动态输入模型导出ONNX实战指南-北京尧图网络

1. 为什么非做“多参数动态输入”导出不可1.1 真实业务里绕不开的多输入模型最近在搞一个多模态内容理解的部署项目训练好的PyTorch模型需要导出成ONNX格式接线上服务。模型本身不算复杂但输入不止一个左边是变长的文本序列右边是任意分辨率的图像batch size还得跟着线上流量走。刚接到需求时我以为就是一行torch.onnx.export的事真正落地才发现“多参数输入动态shape”这套东西里面藏着不少门道尤其是当你以为导出成功、结果跑到线上却报错的时候那种滋味谁经历谁知道。先说多参数。很多教程和示例代码都是拿单输入模型演示的比如一个图像分类网络输入就一个x导出的时候args传一个tensor就行。但实际业务里模型接收多个输入的情况太常见了推荐系统里用户行为序列可以拆成id、时间戳、行为类型多路特征分别过不同的Embedding表搜索场景里query和doc两个塔各吃各的输入多模态模型更不用说了文本和图像同时进网络。你需要把这些tensor按模型forward函数里参数出现的顺序打包成一个tuple传给torch.onnx.export的args。这个“顺序”两个字后面会反复坑人。1.2 动态输入到底解决了什么痛点动态输入对应的是动态shape核心就解决三件事。第一线上batch size不稳定。请求高峰可能一次塞32条样本低峰只有1条。如果导出时把batch维写死成32线上跑batch1要么被迫给其他31个位置补padding浪费算力和显存要么干脆直接报错。第二序列长度天然不固定。文本有长有短同一个batch里每条样本的seq len也不一样动态维度能让模型在推理阶段真正接受变长序列而不是训练时那种靠mask硬撑的“伪变长”。第三图像尺寸可能随预处理策略变化。同一个模型可能被不同上游服务喂不同分辨率的图一张图1920x1080下一张可能640x360只做resize当然可以但有损耗某些视觉任务对分辨率比较敏感。ONNX对动态shape的支持手段就是dynamic_axes把这几个维度在模型描述里标记成符号比如batch_size、seq_len、height、width。推理时这些维度不绑定具体数字框架根据实际输入tensor的shape去推断。这个机制是“支持多参数与动态输入”的技术基础也是整个导出环节最需要小心设计的地方。1.3 为什么是ONNX而不是TorchScript既然PyTorch本身提供TorchScript为什么还要折腾ONNX我在这个项目里也认真对比过。TorchScript的好处是跟PyTorch搭配无缝trace或script一下就得到torch.jit.ScriptModule可以直接用libtorch做C服务。但问题在于它的生态基本被PyTorch自己圈住了想再接TensorRT这类推理加速引擎流程反而更绕。ONNX是开放的中间格式几乎所有主流推理框架都拿它当第一站。ONNX Runtime、TensorRT、OpenVINO这些引擎都提供了直接吃ONNX的入口服务端用C、Java、Go都有对应SDK。换句话说ONNX承担了一层“跨框架转换”的职责训练用PyTorch导出成ONNX之后后续转成TensorRT的engine还是在ONNX Runtime上做INT8量化选择空间都很大。这个项目最终走的就是“PyTorch训练 - ONNX导出 - 部署引擎推理”的链路ONNX在这条链路里是承上启下的关键节点绕不开。2. 先搞清楚这几个原理导出时心里才有底2.1 torch.onnx.export默认走的是trace不是“懂你”很多初学者以为torch.onnx.export会“理解”整个模型的前向逻辑其实它默认走的是JIT trace给一组样例输入模型正常跑一次前向PyTorch把这期间真实执行过的算子依次记录下来生成计算图。话句话说trace是一台“录像机”而不是“翻译器”。这个机制带来一个关键结论模型里的Python控制流比如if、for、while在trace时只会记录这次运行实际走到的分支。比如输入长度大于100走分支A否则走分支B而样例输入长度是50那图上就只有分支B的算子线上来了长度120的样本模型也不会“灵机一动”切到分支A它根本没有那个分支。所以导出前有几件事是必须做的。一是model.eval()确保BatchNorm和Dropout切到推理状态否则导出的图带着训练时的随机行为线上结果大概率不对。二是准备一组合适的样例输入shape要能代表线上常见的状态比如batch设2而不是1因为某些算子会把批量维度折叠成常量batch1时更容易触发这个问题。三是尽量让模型forward保持线性结构把控制流逻辑留在模型外部的预处理/后处理环节里。2.2 多参数传参tuple顺序比代码注释更重要torch.onnx.export的args参数在多输入模型下必须传一个tuple里面每个元素对应forward里的一个输入顺序绝对不能乱。比如模型定义是def forward(self, token_ids, image_input)那args就要写成(dummy_token_ids, dummy_image_input)而不是反过来。这里有个容易忽略的细节样例输入要和你线上推理时的tensor类型、layout保持一致。用torch.randint生成token序列时记得指定dtypetorch.long图像用torch.randn生成float32两者都要先detach不设requires_grad。否则trace时计算图里可能残留梯度相关的节点虽然不一定会报错但会让图变得不干净增加排查负担。我习惯在导出前先打印一份输入输出形状的注释写进代码里方便后面核对# token_ids: [B, S] - 长整型 # image_input: [B, 3, H, W] - 浮点型 # logit: [B, 1]看起来像废话但在你对着报错信息抓耳挠腮的时候这段注释能帮你快速确认到底是不是把参数顺序喂反了。2.3 dynamic_axes的两种写法和命名习惯dynamic_axes可以用来指定输入和输出中哪些维度是动态的。官方文档给了两种写法实际项目里我推荐第一种。第一种带维度名的dict写法dynamic_axes{ token_ids: {0: batch_size, 1: seq_len}, image_input: {0: batch_size, 2: height, 3: width}, logit: {0: batch_size}, }第二种只标维度下标的列表写法dynamic_axes{ token_ids: [0, 1], image_input: [0, 2, 3], logit: [0], }第一种写法会在模型文件里留下语义清晰的维度名比如batch_size、seq_len打开Netron看图的时候一眼就能看出哪些维度是可变的下游推理引擎的配置界面里也更容易对应上。第二种写法更省事但可读性差尤其在模型层数多、中间tensor多的时候排查问题非常痛苦。还有一个细节动态维度的名称不要用None、0这种容易被框架当成常量的值也不要把输入动态维度命名成和输出维度一样但含义不同的名字比如输入batch叫batch_size输出batch又换成N虽然不会报错但会让引擎侧的配置变得混乱。统一命名是低成本高收益的好习惯。3. 完整实例一个文本加图像的双输入模型导出3.1 准备一个合理的示例模型理论讲再多不如跑一个能复现的例子。我这里构造一个双输入模型文本序列过Embedding加LSTM图像过一层卷积加全局池化两路特征拼接后再接一个分类层。代码很简单但它具备了多参数和动态输入的所有关键特征。import torch import torch.nn as nn class MultiInputDemo(nn.Module): def __init__(self, vocab_size10000, emb_dim64, hidden_dim128): super().__init__() self.embedding nn.Embedding(vocab_size, emb_dim) self.lstm nn.LSTM(emb_dim, hidden_dim, batch_firstTrue) self.cnn nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d((1, 1)), ) self.image_fc nn.Linear(16, hidden_dim) self.cls nn.Linear(hidden_dim * 2, 1) def forward(self, token_ids, image_input): # token_ids: [B, S] # image_input: [B, 3, H, W] emb self.embedding(token_ids) seq_out, _ self.lstm(emb) # [B, S, hidden] text_feat seq_out[:, -1, :] # 取最后一个时间步 img_feat self.cnn(image_input).flatten(1) # [B, 16] img_feat self.image_fc(img_feat) # [B, hidden] fused torch.cat([text_feat, img_feat], dim-1) return self.cls(fused) # [B, 1] model MultiInputDemo() model.eval()模型本身不重要重点是它涉及了三个关键点两个输入、文本序列变长、图像尺寸可变。图像分支用了AdaptiveAvgPool2d把特征压缩到1x1这样后续的Linear就不依赖输入图的分辨率是为动态H和W“铺路”的必要设计。如果你在模型里写死image.flatten(1)后再接一个输入维度等于3x224x224的线性层那导出的动态宽高注定是废的后面会说为什么。3.2 导出参数逐个拆解准备好样例输入开始导出。样例输入我只在形状上做文章batch给2序列给16图像给3x224x224。batch给2而不是1的原因前面提过这里不再重复。bs, seq_len 2, 16 dummy_token_ids torch.randint(0, 10000, (bs, seq_len), dtypetorch.long) dummy_image_input torch.randn(bs, 3, 224, 224, dtypetorch.float32) torch.onnx.export( model, args(dummy_token_ids, dummy_image_input), fmulti_input_dynamic.onnx, export_paramsTrue, opset_version13, do_constant_foldingTrue, input_names[token_ids, image_input], output_names[logit], dynamic_axes{ token_ids: {0: batch_size, 1: seq_len}, image_input: {0: batch_size, 2: height, 3: width}, logit: {0: batch_size}, }, ) print(exported)逐个解释这几个参数因为每个都有讲究。export_paramsTrue表示把权重参数固化进ONNX文件一般来说线上推理都是固定权重所以这个参数保持默认即可。do_constant_foldingTrue让导出器做常量折叠图上能预先算好的部分就先算掉减少推理时的计算量但它也可能把一些shape相关的常量提前算死后面排查动态维度失效时要留意。opset_version13这里先说结论后面专门讲版本怎么选。input_names和output_names给模型的输入输出起了名字这个名字会写进ONNX文件不依赖你在PyTorch里的变量名。之后用ONNX Runtime推理时feed_dict里的key必须和这里完全一致。dynamic_axes把token_ids的第0维和第1维、image_input的第0维和第2维第3维、logit的第0维标记成动态。3.3 结构校验动态维度到底生效没有导出完别急着部署第一件事是校验文件结构和动态维度是否真正写进图里。我用的是自带的onnx.checker配合shape_inference。import onnx from onnx import shape_inference onnx_model onnx.load(multi_input_dynamic.onnx) onnx_model shape_inference.infer_shapes(onnx_model) onnx.checker.check_model(onnx_model) for inp in onnx_model.graph.input: dims [] for d in inp.type.tensor_type.shape.dim: dims.append(d.dim_param if d.dim_param else d.dim_value) print(finput {inp.name}: {dims}) for out in onnx_model.graph.output: dims [] for d in out.type.tensor_type.shape.dim: dims.append(d.dim_param if d.dim_param else d.dim_value) print(foutput {out.name}: {dims})如果一切正常你会看到类似这样的输出input token_ids: [batch_size, seq_len] input image_input: [batch_size, 3, height, width] output logit: [batch_size, 1]看到batch_size、seq_len这些字符串而不是具体数字说明动态维度已经写进模型描述里了。注意image_input的第1维是3这是通道维我们没有把它动态化符合预期。logit的第1维是1分类输出维度固定也没问题。这一步看似简单却能筛掉很多问题。如果输出里显示的是数字2、16、224说明动态轴定义没有生效或者被中间某个算子固化了直接往下走部署流程就是在埋雷。3.4 端到端验证用ONNX Runtime跑动态shape文件校验通过只能说明结构合法还不能说明推理行为正确。真正有说服力的验证是用ONNX Runtime加载模型换上完全不同的shape跑一遍跟PyTorch原模型的输出做对比。import numpy as np import onnxruntime as ort session ort.InferenceSession( multi_input_dynamic.onnx, providers[CPUExecutionProvider], ) for bs_test, seq_test, hw_test in [(1, 8, 128), (4, 32, 256)]: test_token np.random.randint(0, 10000, (bs_test, seq_test)).astype(np.int64) test_image np.random.randn(bs_test, 3, hw_test, hw_test).astype(np.float32) ort_logit session.run( [logit], {token_ids: test_token, image_input: test_image}, )[0] with torch.no_grad(): torch_logit model( torch.from_numpy(test_token), torch.from_numpy(test_image), ).numpy() diff np.abs(ort_logit - torch_logit).max() print(fbatch{bs_test}, seq{seq_test}, hw{hw_test}, max_abs_diff{diff:.6e})这里的测试shape故意和样例输入不一样batch从2变成1和4序列从16变成8和32图像尺寸从224变成128和256。只要动态轴配置正确ONNX Runtime就能正常吃进去。而max_abs_diff的量级能告诉我们对齐情况——一般在1e-5左右属于正常浮点误差如果差到1e-1甚至更大那就要查模型状态、算子实现或者常量折叠是否引入问题。我自己在这套验证里翻过车有一版模型没加model.eval()就导出静态shape下偶尔能蒙混过关因为BatchNorm在推理时如果遇到和训练分布差异不大的输入误差没那么明显但换成动态shape、换一批不同长度的输入后误差立刻放大到1e-2以上。所以验证脚本里一定要换shape、换数据分布别只拿同一份样例反复测。4. 实操中踩过的坑和排查实录4.1 动态维度被悄悄固化最常见的隐形杀手我最初导出的一个版本dynamic_axes明明写全了shape_inference却显示输入还是固定数字。排查了半天发现问题出在模型内部使用了x.shape[0]或x.size(0)去参与后续计算。trace机制下x.shape[0]会被当做一个具体数值记录进图里比如样例输入的batch是2那图上所有依赖这个值的计算就全按2算batch_size这个动态轴等于没设。典型雷区包括x.view(x.size(0), -1)、torch.flatten(x, start_dim1)配合前面的shape推断、nn.Flatten()、以及把x.shape当条件做判断的逻辑。碰到的案例是模型里有类似x.view(x.size(0), -1)的写法导出时样例batch2view出来的第二维被算成了固定值。换成x.flatten(1)后还是不行因为在trace下flatten的start_dim是常量但展平时的维度大小依然依赖前面tensor的实际shape一样会被固化。那怎么避坑一个原则是模型内部对shape的依赖越少越好。尽量把“取长度”“判断维度”这类逻辑放到模型外部forward里只接受已经处理好shape的tensor。万一遇到改不动逻辑的模型可以考虑用torch.jit.script做script模式的导出它对动态shape的处理能力比trace强但script对PyTorch算子的覆盖范围有限属于带着新问题换旧问题的方案要用得提前评估。4.2 多输入顺序和名字错乱多输入模型的第二个高频问题是导出时把args里的tensor顺序搞反或者input_names跟Python变量名对不上。例如模型forward定义是forward(self, token_ids, image_input)但导出时args(dummy_image_input, dummy_token_ids)。PyTorch不会主动告诉你喂错了它只会照着错误顺序去执行forward表现就是模型内部某个算子shape不匹配报错。如果输入的shape恰好能匹配比如两个输入恰好维度一样那模型会“静默跑错”特征交叉以后整个语义都变了这种错最坑因为不报错但结果全是垃圾。排查方法很朴素看一眼模型定义里forward参数顺序然后把args里的tensor按同样顺序排列。input_names也建议和forward参数名保持一致或至少别相差太远。ONNX Runtime的session.run里feed_dict的key必须和input_names一一对应而且不能写错一个字母否则直接报Invalid Feed Input Name。4.3 opset版本带来的兼容性问题opset是ONNX算子集的版本号版本越高能用的算子和对动态shape、Loop、If等结构的支持越好。但高版本对下游推理引擎的压力也越大。我遇到过导出自带opset 17结构校验通过到了某嵌入式部署环境的推理SDK上却直接报“Unsupported operator”的情况最后排查下来是那套SDK的ONNX parser只支持到opset 11。选opset有几个经验如果走ONNX Runtime做CPU/GPU服务opset 13基本够用兼容性和新特性比较平衡如果之后要转TensorRT建议先查对应TensorRT版本的ONNX parser支持到多少通常11或13比较稳如果模型里用到了比较新的算子比如某些PyTorch原生算子导出的新opset才有的形式那就需要把引擎升级到支持对应版本的而不是用高opset硬顶。另一个相关坑是不同opset下同一个PyTorch算子导出的结构可能不一样。比如ReduceSum在opset 13里keepdims的默认行为有变化Slice的写法也不同于早期版本。所以尽量不要在不同版本之间随便迁移模型文件固定一个opset版本走到底遇到算子错误再单独看是哪个算子在挑版本。4.4 模型内部控制流让trace“失真”前面提过trace只会记录一次执行的路径。如果模型forward里有if seq_len threshold这种判断线上真实的输入可能走到另一个分支而导出的ONNX根本没有那个分支的算子结果自然不对。这种问题的修复思路有三条。第一条是把控制流逻辑前移让模型forward保持纯计算判断和分支在部署代码里做需要走不同分支就预处理出不同的tensor再喂给同一个模型。第二条是如果控制流逻辑简单可以考虑script模式导出PyTorch的torch.jit.script能把部分控制流编译成ONNX支持的Loop/If结构但script对很多Python语法和自定义算子的支持有限适配成本不低。第三条是干脆导出多个模型每个模型固化一个分支路径线上用部署代码选择调用哪一个。这个方案“笨”但稳定尤其适合分支差异很大的模型。4.5 导完但结果对不上的排查路径“能跑但结果不对”是最耗时的坑因为系统不会给你报错。我的排查顺序一般是固定的。第一确认PyTorch这边模型处于model.eval()状态BatchNorm用running mean/varDropout直接关掉这是最大的一类原因。第二固定shape重新导一个ONNX作为基准和动态shape导出的版本做对比如果固定shape也有误差说明问题出在导出本身跟动态轴没关系。第三复现一个最简单的输入比如batch1、短序列、小图像分别跑PyTorch和ONNX Runtime打印每一层输出的shape和值用二分法定位是哪一类算子开始出现差异。第四检查数据流向ONNX Runtime里feed的输入有没有经过和PyTorch前处理一样的归一化。不少“模型部署后结果不对”的项目最后定位到的问题根本不在模型而在前处理代码。浮点误差这个东西要理性看待。不同实现之间存在1e-5到1e-4的差异非常正常如果任务对精度要求不是极端敏感这个量级完全可以接受。但如果出现1e-1级别甚至更大就别用“浮点误差”安慰自己了一定有结构性问题。5. 经验沉淀与后续扩展5.1 我受益最深的一个习惯这一整套折腾下来我最深的体会是动态shape导出不是写几个dynamic_axes参数就完事它考验的是整个模型对shape的敏感度。平时写模型随手写的x.shape[0]、x.view(-1, 3)放到动态shape的导出场景里全都可能变成炸药。如果一个项目从一开始就确定要部署最好在写模型的时候就有意识地把shape相关逻辑隔离起来或者用一个专门的export测试用例跑通动态导出再开始训练别等训练完才回头处理。另外我强烈建议把“导出 - 校验 - 换shape对比”这三个步骤做成一个固定脚本。每次改动模型结构后跑一遍比临时敲命令可靠得多。脚本里保存当时的opset_version、样例输入shape、动态轴定义一旦出问题能快速追溯到是哪一次改动引入的。5.2 从ONNX继续往前走的几个方向ONNX导出只是部署链路的第一站后面还有很大空间。常见的方向包括继续转成TensorRT engine并配置动态shape的optimization profile在ONNX Runtime上做INT8量化降低显存占用和推理延迟或者把ONNX文件裁剪掉一些不需要的输出节点让推理图更精简。接下来如果你的模型还要面对更复杂的部署环境建议先把动态轴命名规范定好然后把验证脚本做成每次CI都跑的检查项。这些基础工作不性感但能让你后面省下大量“救火”的时间。

PyTorch多参数动态输入模型导出ONNX实战指南

相关推荐

统信UOS SSH连接不上？从网络配置到远程维护的完整排查指南

感知与已有知识匹配：WSaiOS-SI结构智能中的知识激活机制

ASP.NET项目中的JavaScript资源管理：依赖顺序与缓存指纹实战

更多资讯

Kubernetes核心概念实战解析：从Pod到Service的协同机制

从HTTP字段到日志审计：网络安全学习第三阶段的流量分析实战

jmeter操作数据库

数据结构.1时间与空间复杂度

超越检索：企业 RAG 的生产级系统设计

IPv4/IPv6寻址翻译器：从NAT64/DNS64原理到实战配置

今日更新

微信云开发旅游小程序实战：LBS打卡+UGC社交全栈方案

金融AI代理平台核心拆解：从多代理协作到财报分析实战

基于大数据的智慧旅游推荐与可视化平台架构实践

本周精选

金融AI落地指南：从调研报告到技术选型与工程实践

大模型上下文成本与注意力优化技术深度解析——为何O(n²)平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测

基于GMM的语音识别系统原理与MATLAB实现详解

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？