ONNX 导出怎样声明动态 batch 和序列长度

ONNX 导出时如果直接使用一组示例输入，导出器通常会把示例里的 batch 和序列长度记录成固定值。想让同一个模型接收不同批大小、不同文本长度，需要显式声明哪些轴可以变化。传统导出路径用 dynamic_axes，而 PyTorch 新的 dynamo=True 导出器优先使用 dynamic_shapes。

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

对文本模型来说，最常见的声明是：输入张量的第 0 轴叫 batch，第 1 轴叫 sequence_length；如果输出也随着文本长度变化，就要为输出的对应轴一起声明动态关系。

动态轴解决的是什么问题

假设示例输入形状是 [1, 128]，不声明动态轴时，ONNX 图会倾向于把它当成固定的 batch=1、序列长度=128。动态轴不会改变张量的阶数，也不会让模型凭空支持任意结构；它只表示某个已有维度的大小在运行时可以变化。

因此，[1, 128] 可以变成 [4, 64] 或 [2, 256]，但仍然是二维输入。模型内部的算子、注意力实现和后端运行时仍可能对可接受范围、显存和最大长度有额外约束。

dynamic_axes 的最小配置

传统 torch.onnx.export 路径中，dynamic_axes 是一个映射：键必须对应 input_names 或 output_names 中的张量名，值可以用“轴编号到轴名称”的字典，或者只列出需要动态化的轴编号。

import torch

# 这是一个示例导出片段：输入形状为 [batch, sequence_length]。
class TextModel(torch.nn.Module):
    def forward(self, input_ids, attention_mask):
        # 示例只保留张量形状关系，真实项目替换为自己的模型计算。
        return input_ids.float().unsqueeze(-1) * attention_mask.float().unsqueeze(-1)

model = TextModel().eval()
input_ids = torch.ones((1, 128), dtype=torch.long)
attention_mask = torch.ones((1, 128), dtype=torch.long)

torch.onnx.export(
    model,
    (input_ids, attention_mask),
    "text_model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    # 第 0 轴是 batch，第 1 轴是文本序列长度。
    dynamic_axes={
        "input_ids": {0: "batch", 1: "sequence_length"},
        "attention_mask": {0: "batch", 1: "sequence_length"},
        "logits": {0: "batch", 1: "sequence_length"},
    },
    # dynamic_axes 适用于传统导出路径；这里显式关闭新导出器。
    dynamo=False,
)

这个配置的关键不是名称本身，而是轴的位置。input_ids 和 attention_mask 都是二维张量，所以第 0 轴和第 1 轴分别对应 batch 与序列长度。示例模型的输出是三维张量 [batch, sequence_length, 1]，因此同样只需要标记前两个轴。

输入、掩码和输出的轴怎样对齐

文本模型通常不只有一个输入。input_ids、attention_mask，以及可能存在的 token_type_ids，前两轴应表达同一组 batch 与序列位置。只给其中一个输入标记动态轴，另一个输入仍可能被导出为固定形状，部署时就会出现输入形状不匹配。

输出是否需要声明，取决于它的形状是否含有可变输入轴。如果输出是每个 token 一个结果，例如 [batch, sequence_length, hidden_size]，就应把输出的前两轴也列入映射；如果输出只是整句分类结果 [batch, num_labels]，通常只把第 0 轴标为 batch。

什么时候改用 dynamic_shapes

PyTorch 官方文档说明，新的基于 torch.export 的 ONNX 导出器在较新的版本中是推荐路径。使用 dynamo=True 时，应优先通过 dynamic_shapes 描述输入形状；dynamic_axes 仍然存在，但属于旧导出逻辑的参数。

import torch

# 新导出器用 Dim 表达可以变化的输入维度。
batch = torch.export.Dim("batch")
sequence_length = torch.export.Dim("sequence_length")

# dynamic_shapes 的嵌套结构要和模型输入结构保持一致。
dynamic_shapes = {
    "input_ids": {0: batch, 1: sequence_length},
    "attention_mask": {0: batch, 1: sequence_length},
}

# 新路径返回 ONNXProgram；保存动作由返回对象负责。
onnx_program = torch.onnx.export(
    model,
    (input_ids, attention_mask),
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_shapes=dynamic_shapes,
    dynamo=True,
)
onnx_program.save("text_model_dynamic.onnx")

两套配置的核心思想相同：把输入的某些维度声明为运行时可变。真正迁移时，应先确认项目使用的 PyTorch 版本和导出器路径，再选择一种写法，不要在同一次调用里把旧示例机械拼到新配置上。

几个容易踩到的边界

• 名称必须对应张量名：dynamic_axes 中写了 ids，但导出参数里实际叫 input_ids，就无法达到预期。
• 动态不等于无限制：序列长度仍受模型位置编码、后端算子、显存和业务最大长度影响。
• 相关输入要同步：掩码和 token 输入的 batch、序列长度必须能一一对应，不能只放宽其中一个。
• 输出按真实形状声明：分类输出常见为二维，不要因为输入是二维就把输出的第 1 轴也误标成 sequence_length。
• 示例输入仍然重要：动态声明只放宽指定轴，示例输入的 dtype、阶数和未声明维度仍会影响导出图。

速查结论

如果你的问题只是“为什么导出的 ONNX 只能接收示例长度”，先看导出器路径，再检查输入和输出名称，最后确认第 0、1 轴是否被正确声明。这样比盲目把所有维度都改成动态更容易定位问题，也更有利于后端优化。

相关问题

dynamic_axes 能不能让输入从二维变成三维？

不能。它只改变已有轴的大小，不能在运行时增加或删除张量维度；模型的输入阶数仍要符合导出图和算子约束。

只有 batch 变化时要怎样写？

只把第 0 轴列为动态即可，例如 "input_ids": {0: "batch"}，序列长度则保留示例中的固定值。

为什么输入放宽了，推理仍然失败？

还要检查其他输入、输出关系和运行时限制，尤其是 attention mask 的形状、未声明的维度、位置编码上限以及后端对动态 shape 的支持范围。