第2章：ONNX模型导出——从PyTorch和TensorFlow说起

好，咱们直接进入正题。ONNX模型导出这件事，说白了就是把训练好的框架模型，转成一种通用的中间格式。为什么要这么干？因为部署的时候，你不可能指望每台设备都装PyTorch或者TensorFlow。ONNX Runtime就是吃这个格式的，所以导出这一步，是后续所有优化的基础。

我个人习惯把导出分成两个维度来看：一是框架本身的API怎么用，二是那些导出参数到底在控制什么。今天咱们就把这两块都讲透。

2.1 从PyTorch导出ONNX模型

PyTorch导出ONNX，核心就一个函数：torch.onnx.export()。我刚开始用的时候，觉得这玩意儿很简单，传个模型、传个输入就完事了。后来踩过坑才发现，参数调不好，导出的模型推理结果直接不对。

先看一个最基础的例子：

import torch
import torchvision.models as models

# 加载一个预训练的ResNet18
model = models.resnet18(pretrained=True)
model.eval()

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

# 导出ONNX
torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 模型的输入张量
    "resnet18.onnx",     # 输出文件名
    export_params=True,  # 是否导出模型参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 是否进行常量折叠优化
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    dynamic_axes=None          # 动态轴设置
)

嗯，这里要注意：model.eval()这一步不能省。我见过有人忘记调成eval模式，结果导出的模型里还带着Dropout和BatchNorm的训练行为，推理结果一塌糊涂。

2.2 从TensorFlow导出ONNX模型

TensorFlow这边稍微麻烦一点。因为TensorFlow的模型保存格式有好几种：SavedModel、H5、Checkpoint等等。我个人推荐用SavedModel格式，因为它最完整，也最不容易出问题。

转换工具用的是tf2onnx，一个独立的Python包。安装很简单：

pip install tf2onnx

然后命令行就能搞定：

python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 11 \
    --inputs input:0 \
    --outputs output:0

当然，你也可以在Python代码里调用：

import tf2onnx
import tensorflow as tf

# 加载SavedModel
model = tf.saved_model.load('./saved_model_dir')

# 获取具体的推理函数
signature = model.signatures['serving_default']

# 指定输入输出的名称
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 转换
output_path = "./model.onnx"
model_proto, _ = tf2onnx.convert.from_function(
    signature,
    input_signature=spec,
    output_path=output_path,
    opset=11
)

这里有个坑：TensorFlow的输入输出名称，你得去模型里查一下。我一般用model.signatures['serving_default'].inputs和.outputs来确认。曾经有一次我写错了名字，导出的ONNX模型输入输出全乱了，排查了半天。

2.3 导出参数详解：input_names, output_names, dynamic_axes

这三个参数，是导出时最核心的配置。咱们一个一个说。

2.3.1 input_names 和 output_names

这两个参数，说白了就是给模型的输入输出张量起个名字。为什么要有名字？因为ONNX模型是一个计算图，每个节点都有名字。推理引擎（比如ONNX Runtime）通过名字来绑定输入输出数据。

举个例子：

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input_1', 'input_2'],  # 如果有多个输入
    output_names=['output_1', 'output_2'] # 如果有多个输出
)

我个人习惯用有语义的名字，比如'image'、'text'、'logits'、'boxes'。这样后续在推理代码里，一眼就能看出每个张量是干什么的。

2.3.2 dynamic_axes：动态轴设置

这个参数，是很多新手容易忽略的。它决定了导出的ONNX模型，是否支持可变大小的输入。

默认情况下，dynamic_axes=None，所有维度都是固定的。比如你传了一个(1, 3, 224, 224)的输入，导出的模型就只能处理batch size为1、图片大小为224x224的数据。

但实际部署中，你往往需要处理不同大小的输入。比如目标检测模型，输入图片可能是各种分辨率。这时候就需要设置动态轴。

看个例子：

torch.onnx.export(
    model,
    dummy_input,
    "model_dynamic.onnx",
    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'、'height'、'width'、'sequence_length'这些。

2.4 导出后的验证

模型导出完了，别急着部署。先验证一下。

我一般用ONNX Runtime来跑一遍推理，和原始框架的结果对比：

import onnxruntime as ort
import numpy as np

# 加载ONNX模型
session = ort.InferenceSession("model.onnx")

# 构造输入
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = session.run(None, {input_name: input_data})

# 和PyTorch结果对比
with torch.no_grad():
    torch_output = model(torch.from_numpy(input_data))

# 检查误差
diff = np.abs(outputs[0] - torch_output.numpy()).max()
print(f"最大误差: {diff}")

如果误差在1e-5以内，基本没问题。如果误差很大，那就要检查是不是某些算子不支持，或者精度设置有问题。

好了，这一章的内容就这些。下一章咱们聊聊ONNX模型的静态优化，包括图优化、常量折叠、算子融合这些实战技巧。到时候我会分享一些我在项目中踩过的坑，保证让你少走弯路。