第4章 ONNX模型导出实战：PyTorch转ONNX、动态轴与静态轴设置、算子兼容性检查

ONNX 这东西，说白了就是模型界的「通用语言」。你拿 PyTorch 训好的模型，想跑到 TensorRT、OpenVINO 或者 ONNX Runtime 上，都得先过 ONNX 这道坎。我刚开始接触部署时，觉得导出 ONNX 不就是一行 torch.onnx.export 的事吗？后来被坑了几次才明白——这里面的门道，比你想象的多得多。

这一章，咱们就手把手把 PyTorch 模型转成 ONNX，把动态轴、静态轴这些概念彻底捋清楚，再聊聊算子兼容性那些坑。嗯，都是实战里每天要面对的东西。

4.1 PyTorch 转 ONNX：核心操作

先来一段最基础的导出代码。假设你有一个训练好的 ResNet-18 分类模型：

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    opset_version=17
)

你看，核心就三个东西：模型、一个 dummy 输入、输出路径。但我个人习惯，一定会加上 input_names 和 output_names。为什么？因为后续用 ONNX Runtime 或者 TensorRT 的时候，你靠名字来绑定数据，而不是靠位置。没名字的话，调试起来特别痛苦。

4.2 动态轴与静态轴：到底怎么选？

静态轴，就是模型导出时输入的形状完全固定。比如上面那个例子，batch size 永远是 1，图片永远是 224x224。动态轴则允许你在推理时改变某个维度的大小——最常见的就是 batch size 和图片宽高。

我遇到过这样一个场景：线上服务需要同时处理不同数量的请求，如果 batch size 固定为 1，吞吐量上不去；如果固定为 8，又可能浪费资源。这时候动态 batch 就派上用场了。

看代码：

torch.onnx.export(
    model,
    dummy_input,
    "resnet18_dynamic.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    },
    opset_version=17
)

这里 dynamic_axes 是个字典，key 是张量名字，value 里指定哪些维度是动态的。我给第 0 维起了个名字叫 "batch_size"，第 2、3 维叫 "height" 和 "width"。这样导出后，ONNX 模型就允许你在推理时传入不同大小的输入。

4.3 算子兼容性检查：别等部署了才发现

模型导出成功 ≠ 模型能跑。最怕的就是导出时没报错，放到 ONNX Runtime 或者 TensorRT 里直接崩掉。原因多半是算子不兼容。

我建议你导出后立刻做两件事：

• 用 ONNX Runtime 跑一遍推理，验证数值是否正确。
• 用 polygraphy 或 onnx.checker 检查模型结构。

先来个简单的检查：

import onnx
import onnxruntime as ort

# 1. 检查模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型结构检查通过 ✅")

# 2. 用 ONNX Runtime 推理验证
session = ort.InferenceSession("resnet18.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

import numpy as np
dummy = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = session.run([output_name], {input_name: dummy})
print("ONNX Runtime 推理成功 ✅ 输出形状:", result[0].shape)

如果这一步报错，最常见的原因是某些 PyTorch 算子 ONNX 不支持。比如 torch.nn.functional.grid_sample 在旧版 opset 里就没有。怎么办？

• 升级 opset（但别太激进）
• 用 torch.onnx.register_custom_op 注册自定义算子
• 改写模型，用更通用的算子替代（比如把 grid_sample 换成双线性插值 + 仿射变换）

4.4 常见算子兼容性速查表

下面这张表是我平时经常参考的，整理了几个容易出问题的算子：

4.5 完整导出流程总结

最后，我把整个流程串起来，给你一个可以直接用的模板：

import torch
import onnx
import onnxruntime as ort
import numpy as np

def export_to_onnx(model, dummy_input, save_path, dynamic=False):
    model.eval()
    # 动态轴配置
    dynamic_axes = None
    if dynamic:
        dynamic_axes = {
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        }
    # 导出
    torch.onnx.export(
        model,
        dummy_input,
        save_path,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes=dynamic_axes,
        opset_version=17,
        verbose=False
    )
    # 验证
    onnx_model = onnx.load(save_path)
    onnx.checker.check_model(onnx_model)
    # 用 ONNX Runtime 跑一次
    session = ort.InferenceSession(save_path)
    dummy_np = dummy_input.numpy().astype(np.float32)
    out = session.run(None, {"input": dummy_np})
    print(f"导出成功，输出形状: {out[0].shape}")
    return onnx_model

# 使用示例
model = torchvision.models.resnet18(pretrained=True)
dummy = torch.randn(1, 3, 224, 224)
export_to_onnx(model, dummy, "resnet18.onnx", dynamic=True)

嗯，这个模板我几乎每个项目都会用。你把它保存下来，以后导出模型直接改改模型和输入就行。