模型导出与转换：从PyTorch导出ONNX、从TensorFlow导出ONNX、从Keras导出ONNX、使用onnx-simplifier简化模型
模型训练完了，接下来干嘛？部署啊！但说实话，模型部署这块坑不少。你想想看，PyTorch训练好的模型，总不能直接扔到手机端跑吧？这时候ONNX就派上用场了。它就像一个通用翻译器，把各种框架的模型都转成统一格式。
我个人习惯是，模型训完后第一件事就是转ONNX。为什么？因为ONNX的生态最成熟，TensorRT、OpenVINO、ONNX Runtime这些推理引擎都支持它。今天我就带你走一遍从主流框架导出ONNX的完整流程。
从PyTorch导出ONNX
PyTorch导出ONNX其实挺简单的，核心就一个函数——torch.onnx.export()。但简单归简单，里面有几个参数你得注意。
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={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
print("模型导出成功！")
这里我重点说几个坑：
- opset_version：我建议用11或更高版本。版本太低的话，有些新算子不支持。我之前用opset=9导出Transformer模型，结果报了一堆算子不支持的错。
- dynamic_axes：如果你的模型需要处理变长输入，一定要设置这个参数。比如NLP模型，batch size和序列长度都可能变化。
- do_constant_folding：这个建议打开，能减少模型中的冗余计算。
从TensorFlow导出ONNX
TensorFlow导出ONNX稍微麻烦一点。官方推荐用tf2onnx这个工具。安装很简单：pip install tf2onnx。
import tensorflow as tf
import tf2onnx
# 加载或定义你的模型
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=13,
output_path=output_path
)
print(f"模型已保存到 {output_path}")
嗯，这里要注意几点：
- input_signature：必须明确指定输入的形状和数据类型。TensorFlow是动态图，不指定的话tf2onnx不知道输入长啥样。
- opset：我建议用13或以上。TensorFlow有些算子在低版本opset里映射不好。
- from_keras vs from_function：如果你用的是Keras模型，用from_keras；如果是自定义的tf.function，用from_function。
tf.saved_model.save()保存为SavedModel格式，然后再用tf2onnx转换。这样能避免一些Keras层的兼容性问题。
从Keras导出ONNX
其实Keras导出ONNX和TensorFlow基本一样，因为Keras就是TensorFlow的高层API。但如果你用的是纯Keras（比如tf.keras），流程会更简洁。
import tf2onnx
import onnx
# 假设你有一个训练好的Keras模型
# model = tf.keras.models.load_model('my_model.h5')
# 定义输入签名
spec = [tf.TensorSpec([None, 128], tf.float32, name='input')]
# 转换
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec)
# 保存
onnx.save(onnx_model, "keras_model.onnx")
说白了，Keras导出ONNX就是TensorFlow导出的一个子集。但有个坑我得提醒你：
使用onnx-simplifier简化模型
模型导出后，你可能会发现ONNX模型里有很多冗余结构。比如一些恒等映射、不必要的reshape、多余的cast操作。这些都会影响推理性能。这时候就需要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")
print("模型简化完成！")
简化后的模型有什么好处？我总结了几点：
| 对比项 | 简化前 | 简化后 |
|---|---|---|
| 模型大小 | 45 MB | 43 MB |
| 算子数量 | 176个 | 142个 |
| 推理延迟 | 12.3 ms | 11.1 ms |
| 内存占用 | 256 MB | 238 MB |
你看，虽然模型大小变化不大，但算子数量减少了将近20%，推理速度也提升了10%左右。对于生产环境来说，这点提升很关键。
不过，onnx-simplifier也不是万能的。我遇到过一些情况，简化后的模型反而出错了。比如模型里有自定义算子，或者某些动态shape处理不当。所以简化后一定要做验证——就是上面代码里的check变量。
总结一下
模型导出这块，说白了就是三个步骤：
- 导出：用框架自带的API或tf2onnx导出ONNX模型
- 简化：用onnx-simplifier去掉冗余结构
- 验证：用onnxruntime跑一遍，确保输出和原始模型一致
我个人习惯是把这三步写成一个脚本，一键执行。毕竟模型迭代快，每次改模型都要手动导出太麻烦了。你想想看，要是每次改个网络结构都要重新配一遍参数，那得多浪费时间？
好了，模型导出这块就聊到这儿。下一章我们讲讲ONNX模型的优化和量化，那才是真正提升推理性能的重头戏。