第4章模型转换基础

第4章：模型转换基础

模型转换，说白了就是把一个框架训练好的模型，搬到另一个框架里跑。这事儿看着简单，但坑不少。我刚开始做端侧部署时，就因为在转换上栽了跟头，折腾了两天才发现是某个算子不兼容。

今天咱们就聊聊ONNX——这个模型转换界的“通用语言”。

4.1 ONNX模型格式介绍

ONNX全称是Open Neural Network Exchange，微软和Facebook联合搞的。它的目标很直接：让不同框架的模型能互相串门。

你想想看，PyTorch训练好的模型，想跑在TensorRT上？或者TensorFlow的模型想部署到Core ML？没有ONNX的话，你得手动重写推理代码，那工作量...嗯，想想就头疼。

ONNX的核心结构其实不复杂：

计算图（Graph）：描述网络的结构，节点是算子，边是张量
算子（Operator）：Conv、ReLU、MatMul这些基础操作
张量（Tensor）：数据流动的载体，包含形状、类型、数值
元数据（Metadata）：模型名称、版本、作者等信息

关键点：ONNX本身不负责训练，它只做推理。所以训练相关的信息（比如优化器状态、梯度）在转换时会被丢掉。

我个人习惯把ONNX看作一个“中间格式”。它不绑定任何硬件，也不偏向任何框架。这种中立性，正是它成为部署标准的原因。

下面这张图展示了ONNX在模型部署流程中的位置：

4.2 PyTorch模型转ONNX

PyTorch转ONNX，官方提供了torch.onnx.export接口。用起来很简单，但细节决定成败。

先看一个基本示例：

import torch
import torch.onnx

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

# 创建一个虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

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

个人经验：opset_version别选太高也别太低。我一般用11或12，兼容性好。太高的话，某些推理引擎可能不支持。

这里有几个坑，我踩过：

动态轴（dynamic_axes）：如果你的模型需要处理不同batch size的输入，一定要设置这个。否则导出的ONNX会固定batch size为1。
模型模式：导出前务必调用model.eval()。我见过有人忘了这步，结果导出的模型里还带着dropout和batch norm的训练行为。
算子兼容性：PyTorch有些操作（比如某些自定义的损失函数）在ONNX里没有对应算子。这时候会报错，需要手动替换。

避坑指南：我曾经在导出时遇到"Unsupported operator"错误，原因是模型里用了torch.einsum。解决办法是把einsum拆成matmul和transpose的组合。

4.3 TensorFlow模型转ONNX

TensorFlow转ONNX，主流工具是tf2onnx。它支持TF 1.x的SavedModel和TF 2.x的Keras模型。

安装很简单：

pip install tf2onnx

转换Keras模型：

import tensorflow as tf
import tf2onnx

# 加载Keras模型
model = tf.keras.models.load_model('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="model.onnx"
)

转换SavedModel：

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --opset 13 \
    --inputs input:0 \
    --outputs output:0

参数	说明	建议
opset	ONNX算子集版本	13或14，兼容性好
inputs/outputs	指定输入输出张量名称	用Netron查看后填写
target	目标部署平台	如tensorrt、cuda等

注意：TF 2.x的模型如果用了@tf.function装饰器，转换时可能会遇到控制流（如tf.while_loop）不兼容的问题。我建议尽量用Keras Sequential或Functional API构建模型。

4.4 ONNX模型可视化

模型转成ONNX后，怎么看它长什么样？推荐两个工具：

Netron：最常用的ONNX可视化工具，支持浏览器在线打开，也支持桌面版
ONNX GraphSurgeon：可以编程方式分析和修改ONNX图

Netron的使用方式：

# 安装桌面版
pip install netron

# 启动
netron model.onnx

或者直接拖拽.onnx文件到 https://netron.app。

用Netron你能看到：

每一层的名称、类型、输入输出形状
参数值（比如卷积核权重）
计算图的拓扑结构

我的习惯：每次转换完ONNX，第一件事就是用Netron打开看看。检查输入输出形状对不对，有没有多余的节点。有一次我发现模型里多了个Identity节点，就是Netron帮我揪出来的。

如果你需要编程方式分析ONNX，可以用ONNX Runtime的Python API：

import onnx

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

# 查看计算图
graph = model.graph
print(f"输入: {[inp.name for inp in graph.input]}")
print(f"输出: {[out.name for out in graph.output]}")
print(f"节点数: {len(graph.node)}")

# 遍历所有节点
for node in graph.node:
    print(f"节点: {node.name}, 类型: {node.op_type}")

嗯，到这里，ONNX的基础知识就差不多了。总结一下今天聊的：

ONNX是模型转换的中间格式，不绑定任何框架
PyTorch转ONNX用torch.onnx.export，注意设置dynamic_axes
TensorFlow转ONNX用tf2onnx，Keras模型转换最省心
可视化用Netron，能帮你快速排查转换问题

下一章咱们聊聊ONNX模型优化和量化，那才是真正让模型在端侧跑快的核心技巧。

📚 课程目录