PyTorch ONNX导出动态轴设置：适配可变输入尺寸
在真实世界的AI应用中，模型很少能以“理想状态”运行。你训练时用的都是固定尺寸的图像或统一长度的文本，但部署到线上后呢？用户上传的照片可能是竖屏480×640，也可能是横屏1920×1080；一段对话可能只有几个词，也可能是一整篇文档。如果模型只能处理预设大小的数据，要么信息被裁剪，要么填充大量无意义的占位符——这不仅浪费计算资源，还可能直接影响预测质量。
于是问题来了：我们能不能让一个从PyTorch导出的ONNX模型，像原始PyTorch模型一样灵活地接受不同尺寸的输入？
答案是肯定的，关键就在于正确配置动态轴（dynamic axes）。
当使用torch.onnx.export()导出模型时，默认情况下，PyTorch会根据传入的dummy_input张量形状将所有维度固化为静态值。比如你传了一个(1, 3, 224, 224)的输入，那导出后的ONNX模型就只会认这个尺寸，哪怕你后续想喂给它(4, 3, 300, 300)，推理引擎也会直接报错。
要打破这种限制，必须显式告诉导出器：“这些维度是可以变的”。这就是dynamic_axes参数的作用。
它接收一个字典结构，键是输入或输出的名字，值是一个子字典，描述哪些维度索引对应什么样的动态语义名称。例如：
dynamic_axes = { "input_tensor": {0: "batch_size", 2: "height", 3: "width"}, "output_logits": {0: "batch_size"} }这里的"input_tensor"是你在导出时指定的输入名，第0维代表批次大小，第2和第3维分别对应高和宽。一旦这样声明，ONNX图中的相应维度就会变成符号化表示（如dim_batch、dim_height），允许运行时动态绑定实际数值。
值得注意的是，这一功能依赖于较新的 ONNX OpSet 版本。建议至少使用opset_version=13，否则某些动态操作可能无法正确转换，尤其是涉及 reshape、transpose 或条件控制流的情况。
来看一个典型的实践案例。假设我们要部署一个用于移动端的图像分类模型，设备来源多样，输入分辨率不一。传统做法是对所有图片 resize 到统一尺寸，但这会导致细节失真或拉伸变形。更好的方式是保留原始比例，由模型内部的自适应池化层（如AdaptiveAvgPool2d）来处理可变空间尺度。
下面这段代码展示了一个简单的CNN网络及其ONNX导出过程：
import torch import torch.onnx as onnx class SimpleCNN(torch.nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv = torch.nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1) self.relu = torch.nn.ReLU() self.pool = torch.nn.AdaptiveAvgPool2d((1, 1)) def forward(self, x): x = self.conv(x) x = self.relu(x) x = self.pool(x) return x.flatten(1) # 实例化模型并设置为评估模式 model = SimpleCNN() model.eval() # 构造示例输入（仅用于推断shape，不影响最终模型灵活性） dummy_input = torch.randn(1, 3, 224, 224) # 指定输入输出名称，便于后续引用 input_names = ["input_tensor"] output_names = ["output_logits"] # 定义动态轴：支持任意batch size及图像高宽变化 dynamic_axes = { "input_tensor": {0: "batch_size", 2: "height", 3: "width"}, "output_logits": {0: "batch_size"} } # 执行导出 onnx.export( model, dummy_input, "simple_cnn_dynamic.onnx", export_params=True, opset_version=13, do_constant_folding=True, input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, verbose=False ) print("✅ ONNX 模型已成功导出，支持动态输入尺寸！")导出完成后，你可以用 ONNX Runtime 加载该模型，并尝试传入不同尺寸的输入进行验证：
import onnxruntime as ort import numpy as np ort_session = ort.InferenceSession("simple_cnn_dynamic.onnx") # 测试多种输入尺寸 for shape in [(1, 3, 224, 224), (2, 3, 300, 400), (1, 3, 128, 128)]: dummy_input_np = np.random.randn(*shape).astype(np.float32) outputs = ort_session.run(None, {"input_tensor": dummy_input_np}) print(f"Input {shape} -> Output {outputs[0].shape}")只要模型结构本身支持可变尺寸运算（比如没有硬编码的 view 操作），就能顺利通过。
当然，灵活性是有代价的。启用动态轴后，部分编译期优化会被削弱。例如TensorRT在构建engine时通常需要明确的最大/最小尺寸作为优化边界；若完全放任维度自由变化，可能导致内存分配策略不够高效，甚至出现OOM风险。因此，在生产环境中常见的做法是：
- 设置合理的动态范围（如 batch ∈ [1, 32], height/width ∈ [128, 1024]）
- 使用 shape inference 工具补全中间节点的维度信息：
from onnx import shape_inference inferred_model = shape_inference.infer_shapes(onnx.load("simple_cnn_dynamic.onnx")) onnx.save(inferred_model, "simple_cnn_inferred.onnx")此外，对于NLP任务，序列长度的动态性尤为重要。BERT类模型常受限于512长度上限，过长则截断，影响语义完整性。通过将input_ids的序列维度设为动态：
dynamic_axes = { "input_ids": {0: "batch", 1: "seq_len"}, "attention_mask": {0: "batch", 1: "seq_len"}, "logits": {0: "batch", 1: "seq_len"} }推理引擎可根据实际文本长度动态调度资源，在保证精度的同时避免不必要的padding开销。
为了确保整个导出流程稳定可靠，推荐结合容器化环境进行标准化操作。比如使用PyTorch-CUDA-v2.9 镜像，它预集成了PyTorch 2.9、CUDA Toolkit、cuDNN以及ONNX相关依赖，开箱即用，极大降低了因本地环境差异导致的导出失败风险。
启动方式如下：
docker run --gpus all -p 8888:8888 -p 2222:2222 pytorch-cuda:v2.9该镜像内置了Jupyter Lab和SSH服务。你可以选择：
- 在浏览器中打开 Jupyter 进行交互式调试，实时查看导出前后模型结构变化；
- 或通过 SSH 登录执行批量脚本，适合集成进CI/CD流水线自动化测试多个模型的导出兼容性。
更重要的是，容器环境天然实现了项目隔离与版本锁定，避免不同项目间PyTorch或ONNX版本冲突的问题。这对于团队协作和长期维护至关重要。
在整个AI部署链条中，动态轴技术扮演着“最后一公里”的角色。训练阶段我们追求泛化能力，而部署阶段则强调鲁棒性和适配性。两者之间的桥梁，正是像ONNX这样的中间表示格式，配合正确的导出配置。
实践中还需注意几点：
- 算子支持度：并非所有PyTorch操作都能完美映射到ONNX，特别是自定义autograd函数或复杂控制流。建议先导出再用
onnx.checker.check_model()验证合法性。 - 命名一致性：
dynamic_axes中的键必须与input_names/output_names完全一致，否则会被忽略且不报错。 - 性能权衡：虽然动态batch有助于灵活响应请求，但不利于GPU满载运行。生产中更常见的是采用动态batching机制（如Triton Inference Server）来平衡延迟与吞吐。
最终你会发现，掌握动态轴设置不仅是技术细节的打磨，更是一种工程思维的体现：如何在灵活性与效率之间找到最佳平衡点。当你不再被“固定输入尺寸”束缚，模型才能真正适应现实世界的复杂多变。
而这一切，只需要几行精心设计的参数配置。
