3、ONNX基础：ONNX格式介绍、PyTorch模型导出ONNX、TensorFlow模型导出ONNX、ONNX模型可视化工具

各位同学，欢迎来到第三章。

这一章我们聊聊ONNX。说实话，ONNX是我在部署工作中打交道最多的中间格式。你想想看，PyTorch训练好的模型，TensorFlow训练好的模型，它们就像两个说不同方言的人。ONNX就是那个翻译官，让它们能在同一个舞台上交流。

3.1 ONNX格式介绍

ONNX全称是Open Neural Network Exchange，说白了就是微软和Facebook牵头搞的一个开放标准。它的核心目标就一个：让不同框架训练的模型能互相转换，最终部署到各种硬件上。

我个人习惯把ONNX看作一个「中间表示层」。它既不是训练框架，也不是推理引擎，它就是个桥梁。嗯，这个比喻很贴切。

ONNX的模型结构其实很简单，就两个核心部分：

• 计算图（Graph）：描述网络的结构，包括节点（算子）和边（张量流动）。
• 权重参数（Initializer）：存储训练好的权重值。

我在项目中遇到过一个问题：ONNX的算子集是有限的。有些PyTorch里的骚操作，比如自定义的F.pad模式，ONNX可能不支持。这时候就得手动替换成ONNX兼容的算子。所以，导出前先检查算子兼容性，能省掉后面一堆麻烦。

3.2 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'}
    }
)

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

• opset_version：我建议用11或更高版本。版本太低，很多新算子不支持。版本太高，有些老旧推理引擎又兼容不了。11是个比较稳妥的选择。
• dynamic_axes：这个参数很实用。如果你需要动态batch size，就把它加上。否则导出的模型只能接受固定大小的输入。
• do_constant_folding：建议开启。它能提前计算一些常量表达式，减少推理时的计算量。

如果你遇到导出失败的情况，别慌。大多数时候是因为模型里有动态控制流（比如if-else、循环）。ONNX要求计算图是静态的，所以这些动态结构需要提前展开。我的做法是：用torch.jit.trace先追踪一遍，看看能不能成功。如果trace都过不了，那ONNX导出基本也没戏。

3.3 TensorFlow模型导出ONNX

TensorFlow导出ONNX，主要靠tf2onnx这个工具。它支持TF1.x的冻结图和TF2.x的SavedModel格式。

先看TF2.x的导出方式：

import tensorflow as tf
import tf2onnx

# 加载一个TF模型
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
)

如果是TF1.x的冻结图（.pb文件），流程稍微复杂一点：

import tf2onnx

# 从冻结图导出
input_names = ["input:0"]
output_names = ["output:0"]

onnx_model = tf2onnx.tfonnx.process_tf_graph(
    tf_graph,
    input_names=input_names,
    output_names=output_names,
    opset=13
)

# 保存为ONNX
with open("model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

这里有个坑，我踩过好几次：TensorFlow的算子命名和ONNX的算子命名不完全一致。比如TF里的FusedBatchNorm，在ONNX里对应的是BatchNormalization。tf2onnx会自动做映射，但偶尔会映射失败。这时候你需要手动指定算子映射表。

3.4 ONNX模型可视化工具

模型导出后，怎么看它长什么样？可视化工具就派上用场了。

我常用的工具有三个：

工具名称	特点	适用场景
Netron	轻量级，支持多种格式，网页版和桌面版都有	日常快速查看模型结构
ONNX GraphSurgeon	可以修改和优化ONNX图，支持Python API	需要手动调整计算图时
ONNX Runtime 自带工具	集成在ORT中，可以打印图信息	调试时快速查看节点信息

Netron是我最常用的。你直接把.onnx文件拖进去，它就能展示出完整的计算图。节点、边、输入输出形状，一目了然。我记得有一次排查一个模型精度问题，就是用Netron发现了一个多余的Reshape节点，删掉后精度就恢复了。

如果你需要编程方式查看图结构，可以用onnx库自带的API：

import onnx

model = onnx.load("resnet18.onnx")
graph = model.graph

# 打印所有节点
for node in graph.node:
    print(f"Node: {node.name}, Op: {node.op_type}")
    print(f"  Inputs: {node.input}")
    print(f"  Outputs: {node.output}")

# 打印所有输入
for input_tensor in graph.input:
    print(f"Input: {input_tensor.name}, Shape: {input_tensor.type.tensor_type.shape}")

# 打印所有输出
for output_tensor in graph.output:
    print(f"Output: {output_tensor.name}, Shape: {output_tensor.type.tensor_type.shape}")

这个方式虽然不如Netron直观，但在自动化脚本里非常有用。比如你想批量检查一批ONNX模型的输入输出是否一致，用代码遍历就比手动一个个看快得多。

好了，这一章的内容就到这里。ONNX作为中间格式，它的核心价值在于「一次导出，多处部署」。下一章我们会深入ONNX的算子集和自定义算子，到时候再聊。