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

好，咱们进入第三讲。前两章我们把环境搭好了，也跑通了第一个推理脚本。但说实话，那个模型还只是个「黑盒子」——你给它输入，它给你输出，中间发生了什么？模型长什么样？能不能换个平台跑？

这一章，我们就来解决这些问题。核心就是 ONNX 这个格式。

3.1 为什么需要ONNX？

先说说我的血泪史。几年前我做项目，用PyTorch训练了一个图像分类模型，效果不错。然后客户说：「我们这边是Java环境，用TensorFlow Serving部署。」我当时就懵了。PyTorch的模型，TensorFlow读不了啊。

你想想看，AI框架这么多——PyTorch、TensorFlow、PaddlePaddle、MXNet……每个框架都有自己的模型格式。就像你用Word写的文档，WPS不一定能完美打开。这就是所谓的「框架锁定」问题。

ONNX（Open Neural Network Exchange） 就是来解决这个问题的。它像一个「通用语言」。你用PyTorch训练，导出成ONNX，然后TensorFlow、ONNX Runtime、甚至手机端都能读。说白了，它就是AI模型的「普通话」。

核心价值：一次导出，到处运行。ONNX让你摆脱框架依赖。

3.2 ONNX格式长什么样？

ONNX文件其实是一个 Protobuf 序列化文件。嗯，这里要注意，它不只是存了权重，还存了完整的计算图。

什么叫计算图？我举个例子。你写了一个简单的神经网络：

输入 -> 卷积 -> ReLU -> 全连接 -> Softmax -> 输出

ONNX会把这一串操作，变成一个「有向无环图」。每个节点是一个算子（比如Conv、Relu），边是张量数据流。这样，任何框架只要读懂这个图，就能复现你的模型。

我个人习惯用 Netron 这个工具来看ONNX图。后面会讲怎么用。

3.3 PyTorch模型导出为ONNX

好，重点来了。怎么把PyTorch模型转成ONNX？

PyTorch官方提供了一个函数：torch.onnx.export()。我直接上代码，你跟着走一遍。

3.3.1 准备一个简单的PyTorch模型

我们先定义一个最简单的卷积神经网络：

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(32 * 26 * 26, 10)  # 假设输入是28x28

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleCNN()
model.eval()  # 记得切换到评估模式
小提示：导出ONNX前，一定要调用 model.eval()。否则BatchNorm和Dropout的行为会不一样，导出的模型推理结果可能不对。我曾经在这上面吃过亏。

3.3.2 构造一个虚拟输入

ONNX导出需要知道输入的形状。我们造一个假数据：

dummy_input = torch.randn(1, 1, 28, 28)  # batch_size=1, 通道=1, 高=28, 宽=28

3.3.3 执行导出

torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 虚拟输入
    "simple_cnn.onnx",   # 输出文件名
    export_params=True,  # 是否导出权重参数
    opset_version=11,    # ONNX算子集版本
    do_constant_folding=True,  # 是否做常量折叠优化
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    dynamic_axes={
        'input': {0: 'batch_size'},   # 允许batch_size动态变化
        'output': {0: 'batch_size'}
    }
)

print("模型导出成功！")

这里我重点说一下 dynamic_axes。默认情况下，ONNX模型是固定batch_size的。但实际部署时，你可能想一次推理多张图片。加上这个参数，你的ONNX模型就支持动态batch了。嗯，这个很实用。

避坑指南：我曾经导出时忘了设 opset_version，结果用了默认的旧版本，导致某些算子不支持。建议至少用 opset_version=11，兼容性最好。

3.4 验证导出的ONNX模型

导出完了，怎么知道对不对？用 onnxruntime 跑一下推理验证：

import onnxruntime as ort
import numpy as np

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

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

# 运行推理
outputs = session.run(["output"], {"input": input_data})

print("ONNX推理结果:", outputs[0].shape)

如果没报错，说明模型导出成功。我建议你每次导出后都跑一下这个验证，确保精度没掉。

3.5 ONNX模型可视化

模型导出来了，但它是「黑盒子」吗？不，我们用 Netron 把它打开看看。

3.5.1 安装Netron

Netron是一个模型可视化工具，支持ONNX、TensorFlow、PyTorch等多种格式。安装很简单：

pip install netron

3.5.2 启动Netron

import netron

# 打开ONNX文件
netron.start("simple_cnn.onnx")

运行后，浏览器会自动打开一个页面。你会看到：

  • 左侧是模型的计算图，每个节点是一个算子
  • 点击节点，右侧显示该算子的详细信息（参数、输入输出形状等）
  • 整个模型的结构一目了然

我记得第一次用Netron看自己的模型时，发现有个卷积层的参数写错了。要不是可视化，我可能调试一整天都找不到问题。

实用技巧：Netron也支持在线版（https://netron.app/），不想装Python也能用。但本地版更安全，适合处理敏感模型。

3.6 常见问题与避坑

最后，我总结几个导出ONNX时容易踩的坑：

问题 原因 解决方案
导出时报错「Unsupported operator」 模型用了ONNX不支持的算子 升级opset_version，或替换为ONNX兼容的算子
推理结果与PyTorch不一致 没有调用model.eval()，或输入数据类型不对 确保eval模式，输入用float32
动态batch不生效 没有设置dynamic_axes 导出时加上dynamic_axes参数
模型文件太大 包含了优化器状态等冗余信息 导出前只保存模型权重，不要保存整个checkpoint

嗯，这一章的内容就到这里。ONNX是AI部署的基石，你掌握了它，后面讲ONNX Runtime、TensorRT都会轻松很多。下一章，我们聊聊如何用ONNX Runtime做高性能推理。

← 上一章 📖 返回目录 下一章 →