ONNX 导出怎么声明动态批次和动态长度

使用当前 PyTorch ONNX 导出器时，优先开启 dynamo=True，再通过 dynamic_shapes 把输入张量的第 0 维声明为动态批次、第 1 维声明为动态长度。两个输入共享同一个 torch.export.Dim 对象，表示它们在运行时必须使用相同的批次和序列长度。

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

现代导出器的最小写法

假设模型的 forward(input_ids, attention_mask) 接收形状为 [batch, sequence] 的两个张量。下面为批次和长度分别创建一个符号维度，并给出合理范围：

import torch

# 同一个 Dim 对象代表多个输入共享同一条形状约束。
batch = torch.export.Dim("batch", min=1, max=64)
sequence = torch.export.Dim("sequence", min=1, max=512)

# 示例输入只用于捕获图，不代表导出后只能接收 2×128。
input_ids = torch.ones((2, 128), dtype=torch.long)
attention_mask = torch.ones((2, 128), dtype=torch.long)

onnx_program = torch.onnx.export(
    model,
    (input_ids, attention_mask),
    input_names=["input_ids", "attention_mask"],
    output_names=["hidden_states"],
    dynamic_shapes={
        # 键对应 forward 参数名，0 是批次轴，1 是序列长度轴。
        "input_ids": {0: batch, 1: sequence},
        "attention_mask": {0: batch, 1: sequence},
    },
    dynamo=True,
)

# 新导出器返回 ONNXProgram，再显式保存模型文件。
onnx_program.save("encoder.onnx")

input_names 给 ONNX 图节点命名，dynamic_shapes 则描述模型输入参数的动态维度。示例中名称保持一致，便于排查，但两者职责不同。

为什么要复用 batch 和 sequence

如果两个输入在业务上必须同形，就不要分别创建 batch_a、batch_b。复用同一个符号维度，会把“input_ids 与 attention_mask 的第 0 维相等”和“第 1 维相等”写进导出约束。这样模型不是接受任意两张二维表，而是接受满足关系的动态输入。

min 与 max 应按部署容量设置。范围过宽可能让后端更难优化，范围过窄则会拒绝合法请求。若模型内部把某个轴固定为常量，导出器可能报告约束冲突；这时要检查模型逻辑，而不是只改大范围掩盖问题。

dynamic_shapes 和 dynamic_axes 怎么选

PyTorch 官方文档明确区分两条路径：dynamic_shapes 为基于 torch.export 的 dynamo 导出器设计；dynamic_axes 是旧 TorchScript 导出方式的参数。不要在新代码中同时堆两套配置来“增加成功率”，先确定导出器，再选择对应参数。

导出失败时先查这四项

• 参数名：dynamic_shapes 字典键要对应 forward 参数名；不确定时可改用与位置参数一致的元组结构。
• 维度关系：必须相等的轴复用同一个 Dim；本来独立的轴不要强行共用。
• 模型约束：reshape、切片、分支或算子可能要求某个维度固定，和动态声明冲突时应按错误提示收紧规则。
• 部署后端：动态 ONNX 只描述符号形状，最终运行时是否支持对应算子、opset 和范围仍要按后端文档确认。

若需要诊断导出过程，dynamo 导出器支持生成报告；但报告是排错材料，不应代替对输入形状关系的设计。最小验收应至少覆盖不同批次、不同序列长度以及上限附近的代表性输入。

常见问题

只把 batch 设为动态可以吗？

可以。只在维度 0 放入 Dim，维度 1 不声明时，长度会按示例或模型约束保持静态。是否需要动态长度取决于真实请求。

输出也要写进 dynamic_shapes 吗？

dynamic_shapes 主要描述模型输入。输出尺寸如果由输入推导，导出图会携带相应符号关系；旧的 dynamic_axes 才常按输入名和输出名分别标记轴。

input_names 能让维度变动态吗？

不能。它只给 ONNX 输入节点命名。动态性来自 dynamic_shapes 或旧路径中的 dynamic_axes。

结语

声明动态批次和动态长度的关键，不是把所有轴都标成“可变”，而是准确表达哪些轴可变、范围多大、哪些输入必须共享同一个符号。新导出链优先用 dynamo=True、dynamic_shapes 与 torch.export.Dim，旧链再保留 dynamic_axes。