ONNX 导出动态维度怎么声明输入轴

ONNX 的动态维度不是把整个输入形状都设成“任意”，而是把指定轴从固定整数改成符号维度。以图像输入 [N, C, H, W] 为例，通常只把 N、H、W 声明为动态，通道轴 C=3 仍保持固定。使用 PyTorch 当前默认的 dynamo=True 导出器时，优先传 dynamic_shapes；只有走 dynamo=False 的旧导出路径时才使用 dynamic_axes。

只需要提前定义好dynamic_axes参数，把输入张量的指定维度索引绑定对应的动态命名，再传入ONNX导出接口，就能完成动态维度的声明，不用修改原有模型的前向逻辑。

PyTorch ONNX 官方文档：https://docs.pytorch.org/docs/stable/onnx.html

ONNX 形状语义：https://onnx.ai/onnx/repo-docs/ShapeAnnotationSemantics.html

先把可变轴与固定轴分开

假设分类模型接收四维图像张量。部署端需要一次处理不同数量的图片，也允许图片高宽变化，那么轴 0、2、3 可以动态；轴 1 表示 RGB 通道，仍应固定为 3。动态维度只改变轴长度，不会改变张量的秩，也不会自动保证模型里的每个算子都支持任意尺寸。

当前导出器用 dynamic_shapes

下面是 dynamo=True 路径的最小写法。dynamic_shapes 的键对应模型 forward 的参数名，这里参数名是 images；内部字典的键是轴索引，值是便于阅读和关联的符号名。

import torch

class Classifier(torch.nn.Module):
    def forward(self, images):
        # 示例模型只演示导出接口，输出保留 batch 维度
        pooled = images.mean(dim=(2, 3))
        return pooled

model = Classifier().eval()
example = torch.randn(2, 3, 224, 224)

# 当前 dynamo 导出器优先使用 dynamic_shapes 声明可变轴
onnx_program = torch.onnx.export(
    model,
    (example,),
    input_names=["images"],
    output_names=["scores"],
    dynamo=True,
    dynamic_shapes={
        "images": {0: "batch", 2: "height", 3: "width"}
    },
)

# 将内存中的 ONNXProgram 保存为模型文件
onnx_program.save("classifier.onnx")

示例输入仍然必不可少，因为导出器要从一次合法调用中取得张量类型、秩和静态维度。dynamic_shapes 只覆盖指定轴；没有列出的通道轴继续使用示例输入中的固定值 3。

旧导出路径才使用 dynamic_axes

很多旧教程使用 dynamic_axes，它本身没有消失，但当前官方文档已经把它和旧导出路径绑定：dynamo=True 时优先用 dynamic_shapes，dynamo=False 时才用 dynamic_axes。旧写法如下：

# 旧导出器通过输入输出名称匹配 dynamic_axes
torch.onnx.export(
    model,
    (example,),
    "classifier-legacy.onnx",
    input_names=["images"],
    output_names=["scores"],
    dynamo=False,
    dynamic_axes={
        "images": {0: "batch", 2: "height", 3: "width"},
        "scores": {0: "batch"},
    },
)

这里 dynamic_axes 的外层键必须和 input_names、output_names 一致。把键写成 Python 参数名但没有在 input_names 中声明，或者只给输入标动态却遗漏需要同步变化的输出 batch，都会让模型接口和预期不一致。

ONNX 里最终保存的是 dim_param

ONNX 的每个张量维度可以保存固定整数 dim_value，也可以保存符号名 dim_param。相同的符号名表示同一次推理中这些轴应绑定到同一个非负整数。例如输入和输出都使用 batch，就表达两者的批大小一致。

可以直接读取模型输入元数据，确认声明是否写入：

import onnx

model_proto = onnx.load("classifier.onnx")
for value_info in model_proto.graph.input:
    dims = []
    for dim in value_info.type.tensor_type.shape.dim:
        # 优先读取符号维度，否则读取固定整数
        dims.append(dim.dim_param or dim.dim_value)
    print(value_info.name, dims)

预期能看到类似 ['batch', 3, 'height', 'width'] 的结构。这里是在检查模型接口元数据，并不等于已经证明所有内部算子都支持任意高宽。

声明成功后还要换尺寸验证

动态维度是接口契约，也是运行时优化的重要依据，但它不能修复模型代码中的尺寸假设。若模型把某个空间维度写死在 view、位置编码或全连接层中，即使输入元数据出现了 height 和 width，换尺寸推理仍可能失败。

最小验收应至少覆盖两组不同形状，例如 [1, 3, 224, 224] 与 [4, 3, 256, 320]。同时检查三件事：模型能被目标运行时加载；两组输入都能完成推理；输出第一维分别为 1 和 4。若只需要动态 batch，就不要顺手把高宽也设为动态，声明越宽，运行时可利用的静态信息越少，模型内部约束也越难排查。

常见误区

dynamic_shapes 的键用 input_names 还是 forward 参数名？

当前 dynamo=True 路径按模型输入参数结构描述动态形状，因此示例中使用 forward(self, images) 的参数名 images。旧 dynamic_axes 则按 input_names 与 output_names 匹配。

同一个符号名能随便复用吗？

不能。相同 dim_param 表示同一次推理中必须取相同值。只有确实相等的轴才应共享 batch、sequence 等名字。

动态高宽是否意味着任意高宽都可用？

不一定。模型结构、算子、下采样倍数、位置编码和目标运行时都可能附加约束。应按真实部署范围选择至少两组不同尺寸做验证，而不是只检查导出文件是否生成。