第一个ONNX模型：从PyTorch导出最简单的ONNX模型

说实话，很多初学者一上来就问我：「老师，ONNX到底难不难？」

我的回答永远是：导出第一个模型，5分钟就够了。但真正理解背后发生了什么，可能需要一点耐心。

今天我们就从零开始，用PyTorch写一个最简单的神经网络，然后把它导出成ONNX格式。你跟着我做一遍，保证能跑通。

准备工作：环境检查

在开始之前，先确认你的环境里装了这些库：

pip install torch torchvision onnx onnxruntime

我个人习惯用 torch 1.10+ 版本，ONNX 的 opset 版本建议用 11 或 13。为什么？因为这两个版本覆盖了绝大多数算子，兼容性最好。

第一步：定义一个简单的PyTorch模型

我们先写一个只有两层全连接的小网络。别嫌它简单——我在项目中遇到过，很多复杂的模型其实也是从这种基础结构堆叠起来的。

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(20, 5)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleModel()
print(model)

这个模型输入是10维向量，经过一个隐藏层（20个神经元），最后输出5维。说白了，就是一个典型的分类器骨架。

第二步：生成一个示例输入

ONNX导出时需要知道输入的形状和类型。你想想看，它得知道「门」有多宽，才能造出合适的「钥匙」对吧？

dummy_input = torch.randn(1, 10)  # batch_size=1, 特征维度=10

这里 1 代表 batch size。实际部署时你可以改成 4 或 8，但导出时用 1 最保险。

第三步：导出ONNX模型

核心代码就这一行：

torch.onnx.export(
    model,               # 要导出的模型
    dummy_input,         # 示例输入
    "simple_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'}
    }
)

嗯，这里要注意几个参数：

• export_params：设为 True 会把权重一起导出，否则只导出网络结构。我建议一直开着。
• opset_version：我踩过坑——用 opset=9 导出时某些算子不支持，后来统一改成 11 就再没出过问题。
• dynamic_axes：这个很关键。它告诉ONNX：batch size 是可变的。你导出时是1，但推理时可以用4、8、16。

第四步：验证导出的ONNX模型

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

import onnxruntime as ort
import numpy as np

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

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

# 运行推理
outputs = session.run(None, {'input': input_data})

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

如果没报错，恭喜你——第一个ONNX模型诞生了！

你还可以用 onnx 库检查模型结构：

import onnx

onnx_model = onnx.load("simple_model.onnx")
onnx.checker.check_model(onnx_model)
print("模型检查通过！")

常见问题与避坑指南

问题	原因	解决方案
导出时报错 Unsupported operator	某些PyTorch算子ONNX不支持	升级opset版本，或替换为支持的算子
推理结果与PyTorch不一致	模型处于训练模式，或输入类型不匹配	导出前调用 model.eval()，确保输入是 float32
动态batch推理报错	没有设置 dynamic_axes	重新导出，加上 dynamic_axes 参数

总结

今天我们从零开始，用PyTorch写了一个两层全连接网络，然后成功导出为ONNX格式。整个过程不到10行核心代码。

你可能会觉得：「就这？也太简单了吧。」

没错，基础用法就是这么简单。但接下来的章节里，我们会逐步深入——处理动态输入、自定义算子、量化、优化...那些才是真正考验功底的地方。

下一章，我们来聊聊ONNX模型的内部结构——看看那个 .onnx 文件里到底藏了什么秘密。

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