第3章 PyTorch模型导出ONNX

3、PyTorch模型导出ONNX：torch.onnx.export核心参数详解与实战

好，咱们直接进入正题。PyTorch转ONNX，最核心的就是torch.onnx.export这个函数。说实话，我刚开始用的时候也踩过不少坑，总觉得参数太多记不住。今天咱们就把这些参数掰开揉碎了讲清楚。

3.1 最基本的调用方式

先看一个最简单的例子。假设你训练好了一个分类模型，想导出成ONNX：

import torch
import torchvision.models as models

# 加载预训练模型
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()。我有个同事就因为这个，导出的模型推理结果总是不对，查了半天才发现是BatchNorm和Dropout还在训练模式。

3.2 核心参数逐个击破

咱们一个一个来看这些参数到底在干什么。

3.2.1 export_params：要不要带权重？

这个参数默认是True。说白了就是：你导出的ONNX文件里，要不要把模型的权重参数也打包进去？

True：导出完整模型，别人拿到就能直接用
False：只导出网络结构，权重另外保存

我个人习惯在生产环境里用True，省事。但如果你要做模型加密或者分步部署，可能会用到False。

3.2.2 opset_version：算子集的版本号

这个参数特别重要。ONNX每个版本支持的算子不一样。版本太低，有些高级操作不支持；版本太高，部署端可能不兼容。

opset版本	特点	推荐场景
9	基础算子，兼容性最好	老旧设备部署
11	支持更多操作，平衡性好	大多数场景
13	支持动态形状、高级操作	需要动态输入
15+	最新特性，但兼容性差	新框架对接

我曾经在项目里用过opset=9，结果模型里有个torch.split操作，导出时报错说算子不支持。后来换成opset=11就解决了。所以我的建议是：默认用11，除非你有特殊需求。

3.2.3 do_constant_folding：要不要做常量折叠？

这个参数很有意思。它会把模型里一些固定的计算提前算好，比如torch.ones(3,3) * 2这种，导出时直接变成常量。

好处是推理速度会快一点，坏处是ONNX文件会变大。我一般都会打开，毕竟推理性能更重要。

3.2.4 input_names 和 output_names：给节点起名字

这两个参数看似简单，但坑不少。名字起得好，后面调试省一半时间。

我建议的命名规范：

输入：input、image、data
输出：output、logits、prediction
如果有多个输入输出，用input_1、output_2这样

为什么要起名字？因为后面用ONNX Runtime推理时，你要通过名字来获取输入输出张量。名字乱起，自己都找不到。

3.2.5 dynamic_axes：动态形状的终极武器

这个参数是很多人的噩梦。先看个例子：

dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

torch.onnx.export(
    model,
    dummy_input,
    "resnet18_dynamic.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes=dynamic_axes
)

这里的意思是：输入的batch_size、height、width都是可变的，输出的batch_size也是可变的。

为什么会需要这个？你想想看，如果你的模型要处理不同尺寸的图片，或者推理时的batch大小不固定，就必须用动态轴。

但要注意：动态轴会降低推理性能。因为ONNX Runtime需要为每种形状重新优化。所以我的建议是：能用固定形状就用固定形状，实在不行再用动态。

3.3 实战：一个完整的导出流程

好了，理论讲完了，咱们来个实战。假设我要导出一个BERT文本分类模型：

import torch
from transformers import BertModel, BertTokenizer

# 加载模型和分词器
model = BertModel.from_pretrained('bert-base-uncased')
model.eval()

# 创建虚拟输入
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text = "Hello, how are you?"
inputs = tokenizer(text, return_tensors="pt")

# 导出ONNX
torch.onnx.export(
    model,
    (inputs['input_ids'], inputs['attention_mask']),
    "bert_model.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input_ids', 'attention_mask'],
    output_names=['last_hidden_state', 'pooler_output'],
    dynamic_axes={
        'input_ids': {0: 'batch_size', 1: 'sequence_length'},
        'attention_mask': {0: 'batch_size', 1: 'sequence_length'},
        'last_hidden_state': {0: 'batch_size', 1: 'sequence_length'},
        'pooler_output': {0: 'batch_size'}
    }
)

print("模型导出成功！")

这里有几个关键点：

BERT有两个输入：input_ids和attention_mask，都要设置动态轴
sequence_length也设成动态，因为不同句子的长度不一样
输出有两个：last_hidden_state和pooler_output

3.4 避坑指南

我曾经踩过的坑，你千万别再踩：

模型没设eval模式：导出前一定要model.eval()，否则BatchNorm和Dropout的行为会变
输入类型不对：ONNX要求输入是torch.Tensor，别传numpy数组
动态轴命名冲突：不同轴的动态名称不能重复，比如{0: 'batch'}和{1: 'batch'}会报错
算子不支持：有些PyTorch操作ONNX没有对应算子，需要手动替换或降级

我的调试小技巧：

导出后一定要验证一下。用ONNX Runtime加载模型，输入同样的数据，对比PyTorch和ONNX的输出差异。差异在1e-5以内就算成功。

3.5 总结

torch.onnx.export的核心参数就这些。记住几个要点：

opset_version用11最稳妥
dynamic_axes能不用就不用
导出后一定要验证
名字起得好，调试没烦恼

下一章咱们聊聊ONNX模型的优化和调试，到时候会用到今天学的这些参数。嗯，先消化消化，有问题随时问我。