第4章 ONNX与模型转换

第4章：ONNX与模型转换

好，咱们今天聊聊ONNX。说实话，我刚入行那会儿，模型转换是个挺头疼的事。PyTorch训练好的模型，想放到树莓派上跑？得，先转成Caffe，再转成TensorFlow，折腾一圈下来，精度还掉了不少。后来ONNX出现了，这玩意儿说白了就是给深度学习模型定了个通用格式，就像USB接口一样，插上就能用。

4.1 ONNX格式介绍

ONNX全称是Open Neural Network Exchange，微软和Facebook联合搞的。它的核心思想很简单：不管你用什么框架训练，最终都转成ONNX，然后部署到任何支持ONNX的设备上。

我个人的习惯是，只要模型要上边缘设备，第一件事就是转ONNX。为什么？因为ONNX Runtime的优化做得太好了，尤其是对ARM架构的芯片。

ONNX的核心优势：

框架互通：PyTorch、TensorFlow、Keras、MXNet都能转
算子标准化：定义了标准化的算子集，避免各家自定义
图优化：自动做常量折叠、节点融合等优化
量化支持：INT8、FP16量化，边缘设备必备

嗯，这里要注意一点：ONNX不是万能的。有些框架特有的算子，ONNX可能不支持。比如PyTorch里的某些自定义操作，转的时候会报错。我遇到过最坑的一次，是用了torchvision里的一个数据增强操作，结果转ONNX时死活过不去。

4.2 PyTorch模型转ONNX

PyTorch转ONNX，官方提供了torch.onnx.export函数。说白了就是三步：定义模型、准备输入、调用导出。

来看个实际例子。假设你训练了一个简单的CNN分类器：

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16 * 15 * 15, 10)
    
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleCNN()
model.eval()  # 记得切到eval模式

# 准备一个虚拟输入
dummy_input = torch.randn(1, 3, 32, 32)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "simple_cnn.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    },
    opset_version=11
)

避坑指南：我曾经在导出时忘了设model.eval()，结果模型里还带着Dropout和BatchNorm的训练状态，推理结果完全不对。记住，导出前一定要切到eval模式。

这里有个关键参数dynamic_axes。你想想看，如果模型只支持固定batch size，部署时多尴尬。设了动态轴，就能支持任意batch size的输入。不过要注意，动态轴会增加一些运行时开销。

4.3 TensorFlow模型转ONNX

TensorFlow转ONNX稍微麻烦点。我建议用tf2onnx这个工具，它比官方自带的转换器更稳定。

安装很简单：

pip install tf2onnx

转换代码：

import tensorflow as tf
import tf2onnx

# 加载你的Keras模型
model = tf.keras.models.load_model('my_model.h5')

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

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

说实话，TensorFlow的算子比PyTorch复杂，尤其是TF2.x里那些自定义层。我遇到过一个问题：模型里用了tf.keras.layers.Rescaling，结果转ONNX时提示不支持。后来我手动把这个层替换成了tf.math.multiply，才搞定。

框架	转换工具	推荐opset版本	常见坑点
PyTorch	torch.onnx.export	11-15	动态轴、eval模式
TensorFlow	tf2onnx	13-15	自定义层、TF2.x算子
Keras	tf2onnx	13-15	Lambda层、自定义损失

4.4 ONNX Runtime部署

模型转好了，怎么跑？ONNX Runtime就是干这个的。它支持CPU、GPU、甚至NPU加速。对于边缘设备，我强烈推荐用ONNX Runtime的C++版本，性能比Python好很多。

先看Python版本的推理代码：

import onnxruntime as ort
import numpy as np

# 创建推理会话
sess = ort.InferenceSession("simple_cnn.onnx")

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

# 准备输入数据
input_data = np.random.randn(1, 3, 32, 32).astype(np.float32)

# 执行推理
output = sess.run([output_name], {input_name: input_data})
print(output[0])

嗯，就这么简单。但实际部署时，有几个优化点要注意：

性能优化建议：

使用ORT的SessionOptions：设置线程数、优化级别
启用图优化：sess.set_optimization_level(ort.GraphOptimizationLevel.ORT_ENABLE_ALL)
INT8量化：边缘设备上，INT8推理速度能快2-3倍
内存复用：避免每次推理都重新分配内存

我曾在RK3588上部署过一个YOLOv5模型。刚开始直接用PyTorch推理，帧率只有15fps。转成ONNX后，用ONNX Runtime跑，帧率到了30fps。再做了INT8量化，直接飙到60fps。你看，这就是ONNX的价值。

4.5 常见问题与调试

模型转换过程中，最常遇到的就是算子不支持。我的调试思路是这样的：

查看错误日志：ONNX会明确告诉你哪个算子不支持
替换算子：用ONNX支持的算子替代，比如用Gemm替代MatMul
降低opset版本：有些新算子老版本不支持，降级试试
使用onnx-simplifier：这个工具能简化模型结构，去掉冗余节点

# 安装onnx-simplifier
pip install onnx-simplifier

# 简化模型
python -m onnxsim input.onnx output.onnx

我曾经遇到一个特别诡异的问题：模型转ONNX后，推理结果全是NaN。查了半天，发现是BatchNorm层在推理时没有正确冻结。解决办法是在导出前手动设置model.eval()，并且把BatchNorm的track_running_stats设为False。

我的调试三板斧：

先用onnx.checker.check_model检查模型完整性
再用onnxruntime跑一遍，对比原始框架的输出
最后用onnx.shape_inference.infer_shapes检查shape是否正确

好了，这一章的内容就这些。ONNX说白了就是个中间格式，但它是连接训练和部署的桥梁。你想想看，没有它，我们得为每个框架写一套部署代码，那得多累。下一章咱们聊聊模型量化，这可是边缘设备上提升性能的杀手锏。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录