ONNX动态形状表示：动态维度符号、Symbolic Shape Inference 与动态轴定义

在ONNX中，“动态形状”并不是一个模糊的概念，而是由一套精确的符号机制来描述的。理解这套机制，是后续处理动态形状问题的地基。本节将深入剖析ONNX如何用动态维度符号（如 batch、seq）表示未知的维度大小，解释Symbolic Shape Inference（符号形状推断）如何在不执行模型的情况下推导张量形状，并厘清动态轴的定义方式及其在模型图中的实际表现。

核心原理：从静态到动态的维度声明

ONNX模型通过 ModelProto 中的 GraphProto 定义计算图，而图的输入输出信息存放在 ValueInfoProto 中。每个 ValueInfoProto 包含一个 TypeProto，其中 TensorShapeProto 用于描述张量的形状。静态形状下，每个维度是一个确定的整数；动态形状下，维度可以是以下三种形式之一：

• 完全未知（未知秩）：不设置 dim 字段，表示该张量的维度数量（秩）都不确定，极少用于输入，多见于中间结果。
• 未知维度大小（动态轴）：维度数量确定，但某一维或某几维的大小不指定，用空 dim 表示，或使用 dim_param 赋予一个符号名称。
• 符号维度（带名称的动态轴）：使用 dim_param 字段，例如 "batch"、"seq"，表示该维度的大小由符号变量决定，且同一符号在模型内多次出现时，其值必须一致。

下面是一个典型的动态形状输入定义（以BERT类模型的 input_ids 为例）：

input {
  name: "input_ids"
  type {
    tensor_type {
      elem_type: INT64
      shape {
        dim { dim_param: "batch" }   // 动态轴：batch size 可变
        dim { dim_param: "seq" }     // 动态轴：序列长度可变
      }
    }
  }
}

这里 batch 和 seq 就是动态维度符号。它们不是占位符字符串，而是参与形状推断的代数变量。当模型图中另一个张量的形状也包含 dim_param: "batch" 时，ONNX运行时和推理引擎会认为这两个维度是绑定的——它们必须相等。

Symbolic Shape Inference：符号形状推断机制

ONNX的Symbolic Shape Inference（符号形状推断）是静态分析引擎的核心能力。它不执行算子，而是沿着计算图的数据流，利用每个算子的形状推断规则，将输入张量的形状信息（包括符号维度）传播到输出张量。这个过程是编译期或加载期完成的，目的是在模型运行前就确定所有中间张量的形状约束，从而为内存分配、算子融合和优化提供依据。

符号形状推断的关键在于：它维护一个符号环境（symbolic environment），将每个 dim_param 映射为一个符号变量。当遇到一个算子时，推断规则会描述输出形状如何由输入形状计算而来。例如：

• Gather：输出形状 = 输入数据形状的前若干维 + 索引形状 + 输入数据形状的剩余维。若索引是动态的，则输出对应维也是动态的。
• Reshape：允许将形状中的 -1 推断为“由总元素数和已知维度推导出的值”。若输入形状包含符号维度，输出形状可能引入新的符号表达式。
• MatMul：要求两个输入的最后两维满足矩阵乘法约束，输出形状的批处理维取两者广播结果，符号维度会进行统一或扩展。

下面通过一个简化的流程示意，展示符号形状推断如何工作：

输入: input_ids [batch, seq]
      attention_mask [batch, seq]
      token_type_ids [batch, seq]

节点1: Shape(input_ids) -> 输出: [2] (静态秩)
节点2: Gather(Shape(input_ids), indices=[0]) -> 输出: [1] 包含符号 batch
节点3: Unsqueeze(节点2, axes=[0]) -> 输出: [1, 1] 包含符号 batch
节点4: Concat(节点3, 其他形状张量) -> 输出: [num_axes] 包含多个符号

在这个流程中，Shape 算子返回输入张量的形状向量，其元素本身是标量张量。当输入维度是符号时，Shape 的输出元素在符号推断中会被标记为“符号值”，而不是具体整数。后续的 Gather、Concat 等操作会将这些符号值组合成新的形状向量，最终用于 Reshape 或 Expand 等需要形状参数的算子。

动态轴的定义与绑定规则

动态轴的定义有两种方式，它们在语义上有细微差别：

在实际工程中，强烈建议使用符号动态轴而非匿名动态轴。原因在于：匿名动态轴虽然也能表示“可变”，但ONNX运行时无法判断两个匿名轴是否应该相等。例如，一个Add算子的两个输入形状分别为 [None, 10] 和 [None, 10]，如果两个 None 是匿名轴，运行时无法保证它们相等，可能导致形状不匹配错误。而如果都命名为 batch，则约束关系一目了然。

案例：Transformer 模型中动态轴的符号传播

以GPT-2或BERT的ONNX导出为例，其核心自注意力模块涉及多个动态轴的交互。假设输入张量 hidden_states 的形状为 [batch, seq, hidden]，其中 batch 和 seq 是符号维度。在多头注意力中，需要将 hidden_states 拆分为多个头：

// 伪代码：Reshape 到 [batch, seq, num_heads, head_dim]
// 然后 Transpose 到 [batch, num_heads, seq, head_dim]

Reshape(hidden_states, shape=[batch, seq, num_heads, head_dim])
  -> 输出形状: [batch, seq, num_heads, head_dim]

Transpose(perm=[0, 2, 1, 3])
  -> 输出形状: [batch, num_heads, seq, head_dim]

这里 num_heads 和 head_dim 通常是静态值（如12和64），而 batch 和 seq 保持为符号。符号形状推断会正确地将 Transpose 后的形状标记为 [batch, num_heads, seq, head_dim]。随后计算注意力分数时，Q和K的点积形状为 [batch, num_heads, seq, seq]，其中两个 seq 都来自同一个符号，因此运行时能确认它们相等，无需显式检查。

如果导出时错误地将 seq 维度设为匿名动态轴，那么Q和K的序列长度维度虽然数值上可能相同，但ONNX运行时无法静态验证，某些优化（如算子融合）会被禁用，甚至在某些推理引擎上直接报错。

误区与复盘：动态形状表示的常见陷阱

下面用一张图总结动态形状表示的核心概念及其关系：

上图展示了三个核心概念如何协同工作：左侧的动态维度符号为形状提供命名变量；中间的符号形状推断利用这些变量沿计算图传播形状信息；右侧的动态轴定义则是在模型接口层面声明哪些维度是可变的。三者共同构成了ONNX处理动态形状的完整机制。

在实际操作中，当你使用 torch.onnx.export 导出动态模型时，需要显式指定 dynamic_axes 参数，其本质就是告诉ONNX导出器：将哪些维度标记为 dim_param，以及使用什么符号名称。例如：

torch.onnx.export(
    model,
    (input_ids, attention_mask),
    "model.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch", 1: "seq"},
        "attention_mask": {0: "batch", 1: "seq"},
        "logits": {0: "batch", 1: "seq"}
    },
    opset_version=17
)

这段代码将 input_ids 和 attention_mask 的第0维命名为 batch，第1维命名为 seq，从而在导出的ONNX模型中生成对应的 dim_param 声明。理解这一底层表示，是后续进行动态形状调试、优化和跨框架部署的前提。