第3章 ONNX基础

3、ONNX基础：ONNX格式介绍、PyTorch模型导出为ONNX、ONNX模型可视化与调试

好，咱们进入第三章。这一章聊的是ONNX，一个在模型部署里绕不开的中间格式。

说实话，我刚开始接触部署时，对ONNX也是一头雾水。心想：我PyTorch模型跑得好好的，干嘛非要转一手？后来踩了几个坑才明白——ONNX就是那个让模型在不同框架、不同硬件之间自由穿梭的“通用语言”。

3.1 ONNX格式介绍：它到底是什么？

ONNX，全称Open Neural Network Exchange，开放神经网络交换格式。说白了，它就是一个模型的中转站。

你想想看，PyTorch训练出来的模型是.pt或.pth，TensorFlow的是.pb或.h5。这些格式彼此不认，但ONNX都能吃进去，也能吐出来。这就是它的价值所在。

ONNX的核心优势：

框架互通：PyTorch→ONNX→TensorFlow，反过来也行
硬件适配：TensorRT、OpenVINO、ONNX Runtime都原生支持ONNX
计算图标准化：ONNX定义了一套统一的算子集，比如Conv、Relu、Gemm

我个人习惯把ONNX看作“模型界的PDF”。你写Word文档，最终导出PDF给别人看，不管对方用WPS还是Office，打开都一样。ONNX也是这个道理。

ONNX的存储结构其实很简单。它用Protobuf序列化，里面包含：

计算图结构（节点、边、张量）
模型参数（权重、偏置）
输入输出元信息（名称、数据类型、形状）

嗯，这里要注意：ONNX本身不负责推理加速，它只是定义了一个标准。真正的加速，要靠TensorRT这类引擎去解析ONNX，然后做优化。

3.2 PyTorch模型导出为ONNX：实操环节

这部分我直接上代码。咱们用一个简单的分类模型做演示。

import torch
import torch.nn as nn

# 定义一个简单的CNN
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 * 13 * 13, 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, 28, 28)

# 导出ONNX
torch.onnx.export(
    model,                    # 模型
    dummy_input,              # 输入张量
    "simple_cnn.onnx",        # 输出文件名
    export_params=True,       # 导出参数
    opset_version=11,         # 算子集版本
    input_names=['input'],    # 输入名称
    output_names=['output'],  # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

print("ONNX导出成功！")

这段代码里，有几个地方我特别想强调一下：

我的经验之谈：

eval模式必须开：我见过有人忘了调model.eval()，结果导出的ONNX里带着Dropout和BatchNorm的训练行为，推理结果完全不对。
opset_version别乱选：我建议用11或更高。太老的版本不支持某些算子，太新的版本可能TensorRT还没跟上。
dynamic_axes很有用：如果你需要处理不同batch size的输入，一定要加上这个。否则导出的ONNX只认固定形状。

导出成功后，你会得到一个simple_cnn.onnx文件。它大概几百KB，里面包含了模型的结构和权重。

3.3 ONNX模型可视化与调试：看清你的计算图

模型导出只是第一步。更关键的是，你得确认导出的ONNX对不对。我见过太多人导出后直接扔给TensorRT，结果报错一脸懵。

可视化工具，我推荐两个：

工具	特点	适用场景
Netron	网页版/桌面版，拖拽即用	快速查看模型结构
ONNX GraphSurgeon	Python库，可编程操作	调试和修改计算图

Netron用起来很简单。你打开它的网页，把.onnx文件拖进去，就能看到完整的计算图。每个节点都可以点开看属性，比如卷积的kernel size、stride、padding。

我个人习惯先用Netron扫一眼，确认输入输出对不对，有没有奇怪的算子。

如果发现有问题，比如某个算子TensorRT不支持，那就得用ONNX GraphSurgeon来修了。

import onnx
from onnx import helper, TensorProto

# 加载ONNX模型
model = onnx.load("simple_cnn.onnx")

# 打印输入输出信息
print("输入:", model.graph.input)
print("输出:", model.graph.output)

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

# 修改节点属性（示例：修改卷积的padding）
for node in model.graph.node:
    if node.op_type == "Conv":
        for attr in node.attribute:
            if attr.name == "pads":
                # 把padding从(0,0,0,0)改成(1,1,1,1)
                attr.ints[:] = [1, 1, 1, 1]

# 保存修改后的模型
onnx.save(model, "simple_cnn_modified.onnx")

我曾经踩过的坑：

有一次我导出一个带Reshape操作的模型，Netron里看着没问题，但TensorRT死活报错。后来用GraphSurgeon一查，发现Reshape的shape输入是个常量节点，TensorRT不认识。我手动把这个常量节点替换成ONNX的Constant算子，问题就解决了。

所以我的建议是：导出ONNX后，先用Netron看结构，再用GraphSurgeon做一次“体检”。别嫌麻烦，这一步能省你后面大量调试时间。

3.4 常见问题与避坑指南

最后，我总结几个导出ONNX时容易翻车的地方：

动态控制流：如果你的模型里有if-else或for循环，ONNX可能导不出来。因为ONNX要求计算图是静态的。
不支持的操作：比如torch.einsum、某些高级索引操作。遇到这种情况，要么换实现方式，要么用ONNX自定义算子。
形状推断失败：有时候导出的ONNX里某些张量的形状是unknown。这通常是因为用了动态shape，但没配好dynamic_axes。