4、ONNX基础与导出：ONNX简介、PyTorch模型转ONNX、动态轴设置、opset版本选择

4.1 ONNX到底是什么？

ONNX，全称Open Neural Network Exchange。说白了，它就是一套模型交换的通用格式。

我刚开始接触深度学习时，最头疼的就是模型迁移。PyTorch训练好的模型，想放到TensorRT上推理？得重写一遍网络结构。想换个框架试试？又得重新训练。那叫一个折腾。

ONNX的出现，就是为了解决这个问题。它定义了一套统一的计算图表示。不管你是用PyTorch、TensorFlow还是其他框架训练的模型，都能导出成ONNX格式。然后，这个ONNX模型就能被各种推理引擎（比如TensorRT、ONNX Runtime）直接加载。

你可以把它想象成深度学习的「通用语言」。嗯，这个比喻很贴切。

4.2 PyTorch模型转ONNX：基础操作

PyTorch转ONNX，靠的是 torch.onnx.export 这个函数。我直接上代码，你一看就明白。

import torch
import torchvision.models as models

# 1. 加载一个预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造一个 dummy 输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出 ONNX
torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 示例输入
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 是否导出参数
    opset_version=11,    # opset 版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output']
)

print("ONNX 导出完成！")

这段代码跑完，你就能在当前目录下看到一个 resnet18.onnx 文件。就这么简单。

不过，这里有几个坑我得提醒你：

• 模型一定要设成 eval 模式。我见过有人忘了这步，结果导出的 ONNX 里还带着 dropout 和 batch norm 的训练行为，推理结果完全不对。
• dummy_input 的 shape 要和实际输入一致。尤其是 batch size 和通道数，别搞错了。
• export_params=True 是默认值，但最好显式写出来。这样代码更清晰。

4.3 动态轴设置：让模型更灵活

刚才的例子，导出的 ONNX 模型是固定 shape 的。也就是说，你只能用 batch size=1、分辨率 224x224 的输入。这在实际部署中，显然不够用。

比如，你的业务场景里，图片大小可能不一样。或者你想在推理时动态调整 batch size。这时候就需要「动态轴」了。

怎么设置？看代码：

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

dynamic_axes 这个参数，就是用来指定哪些维度是动态的。它的格式是字典：

• key 是输入/输出的名称（要和 input_names、output_names 对应）
• value 又是一个字典，key 是维度索引，value 是维度名称（随便起，但要有意义）

上面这个例子，我让 batch size、高度、宽度都变成了动态的。这样，你就能传入任意 batch size、任意分辨率的图片了。

4.4 opset版本选择：别小看这个数字

opset，全称 Operator Set。它定义了 ONNX 支持哪些算子、每个算子的行为是什么。

每次 ONNX 发布新版本，都会增加一些算子，或者修改已有算子的行为。所以，opset 版本的选择，直接决定了你的模型能不能被正确导出和推理。

我整理了一个表格，方便你参考：

opset 版本	PyTorch 版本	关键变化	我的建议
9	1.2+	基础算子支持	太老了，别用
10	1.3+	新增了一些激活函数	兼容性还行，但也不推荐
11	1.5+	支持了更多动态操作	推荐，稳定且功能齐全
12	1.7+	新增了一些量化相关算子	如果你用量化，可以考虑
13	1.8+	支持了更多高级操作	新项目可以用，但注意兼容性
14+	1.10+	持续更新	除非你有特殊需求，否则没必要追新

我个人习惯，默认用 opset 11。为什么？

• 它足够稳定，几乎所有主流算子都支持了
• TensorRT 对 opset 11 的支持最成熟
• 兼容性好，不会出现「导出成功但推理失败」的尴尬

我曾经踩过一个坑：用 opset 13 导出了一个模型，里面用到了某个新算子。结果 TensorRT 不支持，推理直接报错。后来我换回 opset 11，问题就解决了。所以，别盲目追新。

4.5 验证导出的ONNX模型

导出 ONNX 之后，一定要验证一下。别等到部署时才发现模型有问题。

验证方法很简单：

import onnx
import onnxruntime as ort
import numpy as np

# 1. 加载 ONNX 模型
onnx_model = onnx.load("resnet18_dynamic.onnx")

# 2. 检查模型结构是否正确
onnx.checker.check_model(onnx_model)
print("模型结构检查通过！")

# 3. 用 ONNX Runtime 跑一次推理
ort_session = ort.InferenceSession("resnet18_dynamic.onnx")

# 构造一个随机输入
input_data = np.random.randn(2, 3, 256, 256).astype(np.float32)

# 运行推理
outputs = ort_session.run(
    None,
    {'input': input_data}
)

print(f"推理成功！输出 shape: {outputs[0].shape}")

这段代码做了两件事：

1. 结构检查：确保 ONNX 模型本身没有语法错误
2. 推理验证：用 ONNX Runtime 跑一次，确保模型能正常推理

如果这两步都通过了，那你的 ONNX 模型基本就没问题了。接下来就可以放心地交给 TensorRT 去优化了。

4.6 本章小结

好了，这一章的内容就这些。我们来快速回顾一下：

• ONNX 是模型交换的通用格式，解决了跨框架部署的痛点
• PyTorch 转 ONNX 用 torch.onnx.export，注意 eval 模式和 dummy input
• 动态轴 用 dynamic_axes 参数设置，灵活但会牺牲一点性能
• opset 版本 推荐用 11，稳定且兼容性好
• 验证 是必须的步骤，别跳过

下一章，我们会正式进入 TensorRT 的世界。到时候，你会看到 ONNX 模型是怎么被优化成 TensorRT 引擎的。嗯，那才是真正有意思的部分。