2、动态轴与静态轴：配置dynamic_axes参数处理可变输入尺寸

好，咱们接着聊。上一节我们把一个固定尺寸的模型转成了ONNX，跑通了流程。但说实话，那种固定尺寸的模型在实际项目中能用到的场景，真的不多。

你想想看，一个生产环境里的模型，输入尺寸怎么可能一成不变？图像分类还好说，可以resize到统一大小。但你要是做目标检测、OCR、语音识别，输入长度几乎每次都不一样。这时候，动态轴就派上用场了。

2.1 什么是动态轴？为什么需要它？

简单说，动态轴就是允许ONNX模型在推理时接受不同尺寸的输入。默认情况下，PyTorch导出ONNX时，所有维度都是固定的。比如你导出一个处理224x224图片的模型，那ONNX就认准了这个尺寸。你传一张300x300的图进去，它直接报错。

我在项目中遇到过这种情况：模型在测试集上跑得好好的，一上线就崩。查了半天，发现是线上图片尺寸五花八门，而ONNX模型只认导出时的那个尺寸。嗯，从那以后，我养成了一个习惯——只要不确定输入永远固定，就开动态轴。

2.2 dynamic_axes 参数详解

PyTorch的 torch.onnx.export() 里有个参数叫 dynamic_axes。它的作用就是告诉ONNX：哪些维度是可变的。

它的格式是这样的：

dynamic_axes = {
    'input_name': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output_name': {0: 'batch_size'}
}

解释一下：

• 键是张量的名字（就是你在 input_names 和 output_names 里定义的名字）
• 值是一个字典，键是维度索引（从0开始），值是你给这个维度起的别名（随便起，但要有意义）

我个人习惯把维度别名起得直观一点，比如 batch_size、seq_len、height、width。这样后面看ONNX模型的时候，一眼就能知道哪个维度是干啥的。

2.3 实战：导出带动态batch的模型

咱们拿一个简单的CNN模型来演示。假设输入是 [batch, 3, H, W]，我想让 batch、H、W 都是动态的。

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3)
        self.fc = nn.Linear(16 * 222 * 222, 10)  # 注意这里

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleCNN()
model.eval()

# 创建一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 配置动态轴
dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    dummy_input,
    'dynamic_cnn.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes,
    opset_version=11
)

print("导出成功！")

所以，如果你的模型里有全连接层，那 H、W 就不能是动态的。要么用全局平均池化把特征图压成固定长度，要么只用 batch 维度做动态。

2.4 实战：导出带动态序列长度的Transformer模型

Transformer模型是动态轴最典型的应用场景。序列长度几乎不可能固定。我建议你养成习惯，所有NLP模型导出时都开动态轴。

import torch
from transformers import BertModel

model = BertModel.from_pretrained('bert-base-uncased')
model.eval()

# dummy输入：batch=1, seq_len=128
dummy_input_ids = torch.randint(0, 30522, (1, 128))
dummy_attention_mask = torch.ones(1, 128)

dynamic_axes = {
    'input_ids': {0: 'batch_size', 1: 'seq_len'},
    'attention_mask': {0: 'batch_size', 1: 'seq_len'},
    'last_hidden_state': {0: 'batch_size', 1: 'seq_len'},
    'pooler_output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    (dummy_input_ids, dummy_attention_mask),
    'bert_dynamic.onnx',
    input_names=['input_ids', 'attention_mask'],
    output_names=['last_hidden_state', 'pooler_output'],
    dynamic_axes=dynamic_axes,
    opset_version=11,
    do_constant_folding=True
)

print("BERT动态模型导出成功！")

这里我把 input_ids 和 attention_mask 的 batch 和 seq_len 都设成了动态。输出也是对应的动态维度。这样，你推理时传 (1, 64) 或者 (8, 512) 都没问题。

2.5 验证动态轴是否生效

导出之后怎么确认动态轴真的生效了？我一般用 onnxruntime 跑一下不同尺寸的输入来验证。

import onnxruntime as ort
import numpy as np

# 加载模型
session = ort.InferenceSession('bert_dynamic.onnx')

# 测试不同尺寸
for seq_len in [32, 64, 128]:
    dummy_ids = np.random.randint(0, 30522, (2, seq_len)).astype(np.int64)
    dummy_mask = np.ones((2, seq_len), dtype=np.int64)
    
    outputs = session.run(
        None,
        {'input_ids': dummy_ids, 'attention_mask': dummy_mask}
    )
    print(f"seq_len={seq_len}, output shape: {outputs[0].shape}")

# 输出应该类似：
# seq_len=32, output shape: (2, 32, 768)
# seq_len=64, output shape: (2, 64, 768)
# seq_len=128, output shape: (2, 128, 768)

看到没？输出形状的第二维跟着输入变了。这就是动态轴在起作用。

2.6 避坑指南

我曾经在动态轴上栽过几个跟头，分享给你：

• opset_version 别太低：动态轴在 opset 11 之后才比较稳定。我建议至少用 opset 11，有条件直接上 opset 15 或 17。
• 动态轴不是万能的：有些算子不支持动态形状。比如 Reshape 如果目标形状是固定的，那动态轴就失效了。导出时留意一下警告信息。
• 性能权衡：动态轴会引入额外的形状计算开销。如果你的生产环境输入尺寸非常固定，用静态轴反而更快。我一般只在需要可变尺寸时才开动态轴。
• 名字别重复：不同维度的别名不能一样。比如你不能把 input 的 dim0 和 output 的 dim0 都叫 batch_size，虽然逻辑上它们确实都是 batch 维度，但 ONNX 会报错。

2.7 总结

动态轴是ONNX导出中最实用的功能之一。说白了，它就是让你的模型从「死板」变得「灵活」。我个人建议：除非你有100%的把握输入永远固定，否则一律开动态轴。多花30秒配置一下，能省掉后面几小时的排查时间。

下一节我们会聊ONNX的算子兼容性问题——为什么有些PyTorch模型导出会报错，以及怎么解决。到时候见。