第4章：模型导出与转换：从PyTorch/TensorFlow导出ONNX、onnx-simplifier优化、模型验证

好，咱们进入第四章。这一章，说白了就是解决一个核心问题：怎么把训练好的模型，变成MCU能吃的“干粮”。

你想想看，PyTorch或者TensorFlow训练出来的模型，就像刚出锅的红烧肉——香是香，但没法直接打包带走。ONNX就是那个“真空包装机”，把模型封装成标准格式。而MCU呢，它胃口小，吃不了大块肉，我们还得把肉切成肉末——这就是量化和优化要做的事。

我个人习惯把这一步叫做“模型的中转站”。搞不定这一步，后面在MCU上跑推理就是空谈。我记得刚入行那会儿，花了两周训练了一个不错的分类网络，结果导出ONNX时各种报错，差点把电脑砸了。嗯，咱们今天就把这些坑提前填上。

4.1 从PyTorch导出ONNX

PyTorch导出ONNX，官方提供了torch.onnx.export接口。看起来简单，但里面门道不少。

先看一个最基础的例子：

import torch
import torch.onnx

# 假设我们有一个训练好的模型
model = torch.load('my_model.pth')
model.eval()

# 创建一个 dummy input，形状要和模型输入一致
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,               # 模型
    dummy_input,         # 示例输入
    'my_model.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'}
    }
)

print("ONNX导出完成！")

这里有几个关键点，我重点说一下：

  • opset_version：我个人建议用11或12。太低的版本不支持一些新算子，太高的版本MCU上的推理引擎可能不支持。我在项目中吃过这个亏，用了opset 15，结果在某个老版本ONNX Runtime上直接崩了。
  • dynamic_axes：如果你的模型需要支持动态batch，这里一定要配。但说实话，在MCU场景下，我通常固定batch为1，省心。
  • do_constant_folding：这个建议打开。它会把一些固定的计算提前算好，减小模型体积。
⚠️ 注意：导出前一定要调用 model.eval()。我见过有人忘了这步，结果导出的ONNX里包含了Dropout和BatchNorm的训练行为，推理结果完全不对。

4.2 从TensorFlow导出ONNX

TensorFlow这边，我们通常用tf2onnx这个工具。流程稍微绕一点，但也不复杂。

先安装：

pip install tf2onnx

然后导出：

import tensorflow as tf
import tf2onnx

# 加载SavedModel格式的模型
model = tf.saved_model.load('my_tf_model')

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

# 转换为ONNX
output_path = "my_model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec, 
    opset=11,
    output_path=output_path
)

print(f"转换完成，保存至 {output_path}")

这里有个坑：TensorFlow的模型格式五花八门，有SavedModel、H5、Keras、Checkpoint等。我个人建议统一用SavedModel作为中间格式，兼容性最好。

我曾经遇到一个项目，客户给的是H5格式，结果tf2onnx死活转不过去。后来我先把H5加载成Keras模型，再保存为SavedModel，最后才转成功。嗯，多一层转换，多一份保险。

4.3 onnx-simplifier：给模型“瘦身”

导出的ONNX模型，往往包含很多冗余的算子。比如一些形状推理、恒等映射、不必要的转换等。这些在GPU上跑没问题，但在MCU上就是负担。

这时候就需要onnx-simplifier出场了。它的作用就是简化计算图，去掉那些“虚胖”的部分。

安装：

pip install onnx-simplifier

使用：

import onnx
from onnxsim import simplify

# 加载原始模型
model = onnx.load('my_model.onnx')

# 简化
model_simp, check = simplify(model)

# 检查是否简化成功
assert check, "简化失败！"

# 保存简化后的模型
onnx.save(model_simp, 'my_model_simplified.onnx')
print("简化完成！")

你也可以用命令行：

onnxsim my_model.onnx my_model_simplified.onnx

简化后，模型体积通常能减少10%~30%。我见过最夸张的一个例子，一个ResNet50从200MB简化到了150MB。为什么？因为原始导出时，PyTorch把很多中间计算都保留了下来，简化器把这些没用的节点全剪掉了。

💡 小技巧：简化后一定要用 check 验证。我遇到过简化器把某些关键节点误删的情况，导致输出结果完全错误。所以，简化后的模型必须重新验证。

4.4 模型验证：确保转换没出问题

模型导出和简化之后，最重要的一步来了——验证。这一步绝对不能省。

验证的核心思路很简单：用同一个输入，分别在原始框架（PyTorch/TF）和ONNX上跑推理，对比输出结果的差异。

我一般这样写验证代码：

import onnx
import onnxruntime as ort
import numpy as np
import torch

# 1. 加载ONNX模型
onnx_model = onnx.load('my_model_simplified.onnx')
onnx.checker.check_model(onnx_model)  # 检查模型结构是否合法

# 2. 创建ONNX Runtime推理会话
ort_session = ort.InferenceSession('my_model_simplified.onnx')

# 3. 准备测试数据（用和导出时一样的dummy input）
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 4. ONNX推理
ort_inputs = {ort_session.get_inputs()[0].name: test_input}
ort_output = ort_session.run(None, ort_inputs)[0]

# 5. PyTorch推理（假设原始模型还在）
with torch.no_grad():
    torch_output = model(torch.from_numpy(test_input)).numpy()

# 6. 对比结果
diff = np.abs(ort_output - torch_output).max()
print(f"最大绝对误差：{diff:.6f}")

if diff < 1e-3:
    print("✅ 验证通过！模型转换正确。")
else:
    print("❌ 验证失败！请检查模型导出过程。")

这里有个关键点：误差阈值怎么设？

我个人经验：

  • 对于分类模型，误差在1e-4到1e-3之间通常没问题，因为softmax后取argmax不受影响。
  • 对于回归模型或检测模型，误差要求更严格，最好在1e-5以内。
  • 如果误差超过1e-2，那肯定有问题，别犹豫，回去检查导出过程。
🔑 核心原则：ONNX推理结果和原始框架推理结果，在数值上应该几乎一致。如果差异明显，要么是导出参数配错了，要么是简化器误删了关键节点。

4.5 避坑指南：我踩过的那些坑

做模型导出这么多年，我踩过的坑能写一本书。这里挑几个最常见的：

  1. 算子不支持：有些PyTorch的高级算子（比如torch.einsum、torch.topk的某些模式）在ONNX里没有对应实现。解决办法是：要么换用基础算子重写，要么升级opset版本。
  2. 动态shape问题：如果你的模型输入尺寸不固定，导出时一定要配dynamic_axes。否则ONNX Runtime会报shape不匹配的错误。
  3. 量化模型导出：如果你做了量化感知训练（QAT），导出时要注意。PyTorch的量化模型导出ONNX需要特殊处理，不能直接用torch.onnx.export。
  4. 控制流问题：模型里有if-else或循环？ONNX对控制流的支持有限。我建议在导出前，把这些控制流改成固定的计算路径。

我曾经接手一个项目，模型里用了torch.where做条件判断，结果导出ONNX后，在MCU上推理结果全错。后来我把torch.where改成了mask乘法，问题就解决了。嗯，有时候“笨办法”反而更可靠。

4.6 本章小结

这一章的内容，说白了就是三件事：

  • 导出：把PyTorch/TF模型转成ONNX格式
  • 简化：用onnx-simplifier去掉冗余节点
  • 验证：确保转换后的模型输出正确

这三步走完，你的模型就准备好进入下一章了——量化。量化才是让模型在MCU上跑起来的关键一步，咱们下章见。

记住一句话：导出不验证，后面两行泪。别问我怎么知道的。

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