3、模型导出实战：从PyTorch导出ONNX、从TensorFlow导出ONNX、常见导出问题

模型导出，说白了就是把训练好的模型「翻译」成一种通用语言。ONNX就是这种通用语言。我刚开始做部署时，总觉得导出很简单——不就是torch.onnx.export吗？结果第一次就栽了跟头。嗯，今天咱们就把这块硬骨头啃下来。

3.1 从PyTorch导出ONNX

PyTorch导出ONNX，我个人习惯用torch.onnx.export。这个接口看起来简单，但参数里藏着不少坑。

import torch
import torchvision.models as models

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

# 创建一个 dummy input
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={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有个关键点——opset_version。我建议用11或以上版本。为什么？因为低版本不支持很多新算子。我在项目中遇到过，用opset=9导出Transformer模型，结果一堆算子不支持，折腾了两天。

3.2 从TensorFlow导出ONNX

TensorFlow导出ONNX，路子跟PyTorch不太一样。我一般用tf2onnx这个工具。说白了，它就是个转换器。

方法一：命令行转换

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --opset 13

方法二：Python API转换

import tf2onnx
import tensorflow as tf

# 加载模型
model = tf.saved_model.load('./saved_model')

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

# 转换
output_path = "./model.onnx"
model_proto, _ = tf2onnx.convert.from_function(
    model.signatures['serving_default'],
    input_signature=spec,
    opset=13,
    output_path=output_path
)

你想想看，TensorFlow的图结构比PyTorch复杂。尤其是那些控制流（if、while），转换时容易出问题。我曾经有个模型用了tf.while_loop，转换后ONNX死活跑不对。后来发现是opset版本太低，换成13就好了。

3.3 常见导出问题与解决方案

导出过程中，我踩过的坑能写一本书。这里挑几个最常见的说说。

问题	现象	解决方案
算子不支持	导出时报错「Unsupported operator」	升级opset版本，或替换为ONNX支持的算子
动态形状失败	推理时报形状不匹配	检查dynamic_axes配置，确保输入输出都设置了
精度下降	ONNX推理结果与原始模型不一致	检查是否开启了常量折叠，或某些算子精度不同
内存溢出	导出大模型时OOM	分批导出，或使用onnx-simplifier简化模型

3.4 验证导出结果

导出完别急着用。我建议先验证一下。

import onnx
import onnxruntime as ort
import numpy as np

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

# 2. 运行推理对比
ort_session = ort.InferenceSession("model.onnx")
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = ort_session.run(None, {"input": input_data})

# 3. 与原始模型对比
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data))

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

嗯，这里要注意。误差在1e-5以内算正常。如果差太多，八成是某个算子转换出了问题。

3.5 模型简化与优化

导出的ONNX模型，有时候会包含一些冗余操作。我一般用onnx-simplifier来瘦身。

python -m onnxsim model.onnx model_simplified.onnx

这个工具会做几件事：常量折叠、死节点删除、形状推断。说白了，就是把模型里那些「没用」的节点去掉。我试过，有些模型能缩小30%的体积。

另外，如果模型要在移动端跑，可以考虑量化。ONNX Runtime支持INT8量化，能大幅提升推理速度。

3.6 总结

模型导出这件事，说难不难，说简单也不简单。核心就三点：

• 选对opset版本——11以上是安全牌
• 处理好动态形状——别让batch size卡住你
• 验证精度——导出后一定要跑一遍对比

我见过太多人，导出后直接扔到生产环境，结果线上炸了。嗯，别急，多花10分钟验证，能省下10小时排查时间。

下一章，咱们聊聊ONNX模型的推理优化。到时候我会分享一些实战中的调优技巧，包括如何用ONNX Runtime做性能分析。