📚 课程目录第3章 ONNX模型动态轴设置

3. ONNX模型动态轴设置

任务目标

为ONNX模型设置动态输入尺寸，使其支持可变批量大小（如训练时batch=1，推理时batch=8或16），避免因输入维度固定而无法在不同批量下复用同一模型文件。

环境依赖

需要安装：onnx（≥1.12）、onnxruntime（用于验证）。若已有PyTorch模型，还需torch完成导出；本案例直接构造一个简单ONNX模型演示动态轴设置，不依赖PyTorch。

实现思路

核心操作是修改ONNX图的input维度中的batch轴，将其值设为字符串"batch"（或任意符号名），同时确保后续所有张量中与该输入batch轴关联的维度也被同步符号化。ONNX官方提供了onnx.shape_inference.infer_shapes配合symbolic_shape_infer工具自动完成这一传播，但为展示原理，本案例先手动构造一个含Conv的简单模型，再用官方API完成动态轴设置。

关键完整代码

import onnx
from onnx import helper, TensorProto
import numpy as np

# 1. 构造一个简单模型：输入[N,3,64,64] → Conv(3→16,k=3) → 输出[N,16,62,62]
X = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 64, 64])
Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 16, 62, 62])

conv_w = helper.make_tensor(
    name='conv_w',
    data_type=TensorProto.FLOAT,
    dims=[16, 3, 3, 3],
    vals=np.random.randn(16, 3, 3, 3).astype(np.float32).tobytes(),
    raw=True
)

node = helper.make_node('Conv', ['input', 'conv_w'], ['output'],
                        kernel_shape=[3, 3], pads=[0, 0, 0, 0])

graph = helper.make_graph([node], 'dynamic_batch_demo', [X], [Y], initializer=[conv_w])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
model.ir_version = 8

# 2. 设置动态batch轴
# 方法A：手动将输入和输出的batch维改为符号名
model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = 'batch'
model.graph.output[0].type.tensor_type.shape.dim[0].dim_param = 'batch'

# 方法B（推荐）：使用官方符号形状推断自动传播到中间张量
# 先重置为固定维度再走API
model.graph.input[0].type.tensor_type.shape.dim[0].dim_value = 1
model.graph.output[0].type.tensor_type.shape.dim[0].dim_value = 1
model = onnx.shape_inference.infer_shapes(model, strict_mode=True, data_prop=False)

from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
model = SymbolicShapeInference.infer_shapes(model, auto_merge=True,
                                            guess_output_rank=True)

# 3. 保存与验证
onnx.checker.check_model(model)
onnx.save(model, 'dynamic_batch_model.onnx')

# 验证：用onnxruntime加载并测试不同batch
import onnxruntime as ort
sess = ort.InferenceSession('dynamic_batch_model.onnx', providers=['CPUExecutionProvider'])
for batch in [1, 4, 8]:
    x = np.random.randn(batch, 3, 64, 64).astype(np.float32)
    out = sess.run(None, {'input': x})[0]
    assert out.shape[0] == batch, f"batch={batch} 输出shape错误: {out.shape}"
    print(f"batch={batch} 验证通过, 输出shape={out.shape}")

运行方式

将上述代码保存为set_dynamic_batch.py；
执行python set_dynamic_batch.py；
脚本运行结束后，当前目录生成dynamic_batch_model.onnx文件。

验证方法与判定标准

脚本末尾的循环会自动完成验证：若onnxruntime能分别以batch=1、4、8成功推理且输出第一维与输入一致，则说明动态轴设置成功。控制台打印的三行“batch=N 验证通过”即为通过依据。若在check_model或InferenceSession阶段报错，说明动态轴传播不完整或模型结构非法。

常见问题

中间层维度未符号化：若只改输入输出而不做形状推断，某些算子（如Reshape、MatMul）的中间张量仍带固定batch，运行时可能报错。务必使用SymbolicShapeInference完成全图传播。
opset版本过低：动态轴依赖较新的opset特性，建议使用opset 13及以上。
动态轴名称一致性：同一批次维度在所有张量中应使用相同符号名（如batch），否则部分推理引擎会将其视为不同动态维度。