4、模型导出与ONNX：PyTorch模型导出为ONNX、TensorFlow模型导出、ONNX算子兼容性检查、ONNX-Simplifier优化

好，咱们进入第四章。模型导出，说白了就是把训练好的模型从框架里“解放”出来，让它能在不同的硬件上跑。ONNX就是这个“通用语言”。我刚开始做嵌入式部署时，总觉得这一步很简单，不就是个格式转换吗？结果踩了不少坑。今天咱们就把这些坑一个个填平。

4.1 为什么需要ONNX？

你想想看，PyTorch模型只能在PyTorch环境里跑，TensorFlow模型只能在TensorFlow环境里跑。但嵌入式设备上跑的可能是TensorRT、OpenVINO或者NCNN。怎么办？

ONNX（Open Neural Network Exchange）就是来解决这个问题的。它定义了一套标准的算子集和计算图格式。只要你的模型能导出成ONNX，就能被各种推理引擎加载。我个人习惯，所有模型最终都要过一遍ONNX，哪怕最终部署不用它，也能做个中间检查。

核心思路： 框架无关 → 硬件无关。ONNX是桥梁，不是终点。

4.2 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={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有几个要点，我一个个说。

4.2.1 opset_version 怎么选？

opset_version 决定了ONNX算子集的版本。版本越高，支持的算子越多，但兼容性可能下降。我在项目中遇到过，选opset=11最稳妥，大部分推理引擎都支持。如果你用了比较新的算子（比如Transformer里的某些操作），可能需要opset=13或更高。

我的建议： 先试opset=11，如果报错说某个算子不支持，再逐步升级。别一上来就选最新的。

4.2.2 dynamic_axes 动态轴

这个参数很实用。它允许你指定哪些维度是可变的。比如上面的例子，batch_size可以动态变化。这样导出的ONNX模型就能处理不同batch size的输入。

但要注意，不是所有算子都支持动态轴。有些算子（比如reshape）如果依赖动态形状，可能会导出失败。我曾经被这个坑过——模型在固定batch size下跑得好好的，一改成动态就报错。后来发现是某个自定义算子不支持动态形状。

4.2.3 do_constant_folding 常量折叠

这个选项建议打开。它会将计算图中可以提前计算的常量部分折叠成常数，减少运行时计算量。说白了就是“能算的先算好”。

4.3 TensorFlow模型导出为ONNX

TensorFlow模型导出ONNX，通常用 tf2onnx 这个工具。它支持SavedModel、Keras H5、Checkpoint等多种格式。

先安装：

pip install tf2onnx

然后看一个Keras模型的例子：

import tensorflow as tf
import tf2onnx

# 加载一个Keras模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')

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

# 导出ONNX
output_path = "mobilenetv2.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=11,
    output_path=output_path
)

如果是SavedModel格式，用 from_saved_model 方法：

import tf2onnx

model_proto, _ = tf2onnx.convert.from_saved_model(
    "path/to/saved_model",
    input_names=["input"],
    output_names=["output"],
    opset=11
)

注意： TensorFlow的算子比PyTorch更复杂，尤其是控制流（如tf.while_loop）和自定义层。我遇到过tf2onnx转换失败的情况，通常是因为模型里用了TensorFlow特有的算子，比如tf.image.non_max_suppression。这时候需要手动替换成ONNX支持的算子。

4.4 ONNX算子兼容性检查

模型导出成ONNX后，不代表就能直接用了。你还需要检查算子兼容性。说白了，就是看看ONNX模型里的每个算子，目标推理引擎是否支持。

我常用的工具是 onnxruntime 的 onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph 错误提示。但更系统的方法是使用 onnx 库自带的检查工具。

先安装：

pip install onnx onnxruntime

然后加载并检查模型：

import onnx

# 加载ONNX模型
model = onnx.load("resnet18.onnx")

# 检查模型结构是否合法
onnx.checker.check_model(model)

print("模型检查通过！")

这个 check_model 会检查模型的结构完整性，比如节点连接是否正确、输入输出是否匹配等。但它不检查算子是否被特定推理引擎支持。

要检查算子兼容性，我一般用 onnxruntime 跑一次推理：

import onnxruntime as ort
import numpy as np

# 创建推理会话
session = ort.InferenceSession("resnet18.onnx")

# 获取输入输出信息
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 构造随机输入
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
outputs = session.run([output_name], {input_name: dummy_input})

print("推理成功！输出形状：", outputs[0].shape)

如果这一步没报错，说明ONNX模型本身没问题，且onnxruntime支持所有算子。但如果你要部署到TensorRT或OpenVINO上，还需要用它们各自的工具再做一次兼容性检查。

避坑指南： 我曾经有一个模型，onnxruntime跑得好好的，但转到TensorRT就报错。原因是模型里用了Resize算子，TensorRT只支持特定模式的Resize。所以，最终部署平台才是检验真理的唯一标准。

4.5 ONNX-Simplifier优化

导出的ONNX模型，有时候会包含一些冗余操作。比如PyTorch自动插入的Shape、Gather、Unsqueeze等节点。这些节点在推理时没什么用，反而增加了计算图复杂度。

ONNX-Simplifier就是用来干这个的。它会分析计算图，删除冗余节点，合并可以合并的操作。

安装：

pip install onnx-simplifier

使用：

import onnx
from onnxsim import simplify

# 加载模型
model = onnx.load("resnet18.onnx")

# 简化模型
model_simp, check = simplify(model)

# 检查简化前后是否等价
assert check, "简化后的模型与原始模型不等价！"

# 保存简化后的模型
onnx.save(model_simp, "resnet18_simplified.onnx")

简化后的模型，节点数通常会减少20%-50%。我见过最夸张的一个模型，从500个节点简化到200个，推理速度提升了30%。

小技巧： 简化后一定要用 check 验证等价性。虽然onnx-simplifier很成熟，但偶尔也会出现简化出错的情况。我建议简化后跑一次onnxruntime推理，对比输出结果是否一致。

4.6 总结与最佳实践

好，咱们把这一章的核心点串一下：

  • PyTorch导出： 用 torch.onnx.export，注意opset_version和dynamic_axes的设置。
  • TensorFlow导出： 用 tf2onnx，注意控制流和自定义算子的处理。
  • 兼容性检查： 先用 onnx.checker.check_model 检查结构，再用 onnxruntime 跑一次推理验证。
  • 模型简化： 用 onnx-simplifier 去除冗余节点，提升推理效率。

我个人习惯的流程是：

  1. 从框架导出ONNX（opset=11，打开常量折叠）
  2. 用onnx-simplifier简化
  3. 用onnxruntime验证推理结果
  4. 用目标部署工具（如TensorRT）做最终兼容性检查

这一步走完，你的模型就准备好进入下一章了——量化与精度调优。嗯，咱们下节课见。

← 上一章 📖 返回目录 下一章 →