1、ONNX基础认知：什么是ONNX、为什么需要ONNX、ONNX的生态与定位

大家好，我是你们这趟ONNX实战课的老师。

咱们第一节课，先聊聊ONNX到底是什么。嗯，说白了，它就是一个“模型交换格式”。

你想想看，深度学习框架这么多，PyTorch、TensorFlow、PaddlePaddle……每个框架都有自己的“方言”。你拿PyTorch训练好的模型，想放到TensorFlow里跑推理？那简直像让一个说中文的人和一个说西班牙语的人聊天，鸡同鸭讲。

ONNX就是那个“世界语”。它定义了一套统一的“语法”和“词汇”，让不同框架的模型能互相理解。

1.1 什么是ONNX？

ONNX的全称是Open Neural Network Exchange，开放神经网络交换格式。它由微软和Facebook（现Meta）在2017年联合推出。

我个人的理解，它就是一个“模型描述语言”。你训练好的模型，不管是卷积、池化、还是全连接层，ONNX都用一套标准化的算子（Operator）和计算图（Graph）来描述。

举个例子，PyTorch里的nn.Conv2d，在ONNX里对应的是Conv算子。TensorFlow里的tf.nn.conv2d，最终也会被映射到同一个Conv算子。

核心要点：ONNX不负责训练，也不负责推理。它只负责“描述”模型长什么样。

1.2 为什么需要ONNX？

这个问题，我当年刚入行时也困惑过。为什么大家不直接用PyTorch或者TensorFlow的模型格式呢？

原因其实很简单：生产环境太复杂了。

我在项目中遇到过这样一个场景：

  • 算法团队用PyTorch训练了一个图像分类模型，精度很高。
  • 但部署到服务端时，用的是TensorFlow Serving。
  • 怎么办？要么把PyTorch模型重写成TensorFlow，费时费力还容易出错。
  • 要么，就用ONNX作为中间桥梁。

ONNX解决了三个核心痛点：

  1. 框架解耦：训练和推理可以用不同的框架。你爱用PyTorch训练，用ONNX Runtime推理，完全没问题。
  2. 硬件加速：很多硬件厂商（比如NVIDIA、Intel、ARM）都针对ONNX做了优化。你导出ONNX模型，就能直接利用这些硬件加速库。
  3. 模型互通：团队之间、公司之间交换模型，不用再问“你用的什么框架？”。ONNX就是通用语言。

我的建议：如果你做的项目需要跨团队协作，或者要部署到不同硬件平台，尽早引入ONNX。别等到模型都训好了，才发现格式不兼容，那可就抓瞎了。

1.3 ONNX的生态与定位

ONNX不是孤立的。它背后有一个庞大的生态圈。

咱们来看看这张“生态地图”：

角色 代表项目 作用
训练框架 PyTorch, TensorFlow, PaddlePaddle 导出ONNX模型
推理引擎 ONNX Runtime, TensorRT, OpenVINO 加载并运行ONNX模型
模型优化 onnx-simplifier, onnxoptimizer 简化、优化ONNX模型
硬件厂商 NVIDIA, Intel, AMD, Qualcomm 提供ONNX的硬件加速后端
模型转换 onnx2pytorch, onnx2keras ONNX转其他框架格式

你看，ONNX处于整个模型生命周期的“中间地带”。它不争不抢，但谁都离不开它。

我个人习惯把ONNX比作“模型界的PDF”。你写文档用Word，但发给别人时，最好转成PDF。为什么？因为PDF格式稳定，不管对方用什么软件打开，排版都不会乱。ONNX也一样，它保证了模型结构的“稳定性”和“可移植性”。

注意：ONNX虽然强大，但也不是万能的。有些框架特有的算子（比如PyTorch的torch.where的某些变体），ONNX可能不支持。这时候就需要你手动处理，或者用ONNX的自定义算子扩展。这部分内容，我们后面的章节会详细讲。

1.4 一个简单的例子

光说不练假把式。咱们来看一个最简单的ONNX模型导出示例。

假设你用PyTorch训练了一个简单的线性模型：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

model = SimpleModel()
dummy_input = torch.randn(1, 10)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "simple_model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

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

这段代码做了几件事：

  • 定义了一个简单的全连接网络。
  • 创建一个虚拟输入（dummy_input），用来“跑一次”模型，让ONNX知道计算图长什么样。
  • 调用torch.onnx.export，把模型导出为simple_model.onnx。
  • 指定了输入输出名称，还设置了动态batch（这样你推理时可以用不同大小的batch）。

嗯，这里要注意：dynamic_axes这个参数，我刚开始用的时候经常忘。如果不设置，导出的ONNX模型就固定了batch size为1。你想想看，生产环境里请求量一大，batch size肯定要动态调整的。所以，导出时一定要考虑动态轴。

1.5 避坑指南

最后，分享一个我踩过的坑。

我曾经把一个PyTorch模型导出成ONNX，然后在ONNX Runtime里跑推理。结果发现，推理结果和PyTorch的原始结果对不上！

排查了半天，发现是torch.onnx.export时，opset_version设置得太低了。低版本的ONNX算子集不支持某些操作，导致计算图被“降级”了，精度自然就变了。

我的建议：导出时，尽量使用最新的opset_version。你可以这样查当前ONNX支持的最高版本：

import onnx
print(onnx.defs.onnx_opset_version())

然后导出时指定：

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=onnx.defs.onnx_opset_version()  # 使用最新版本
)

好了，第一节课就到这里。咱们把ONNX的“是什么、为什么、怎么用”理清楚了。下一节课，我会带你深入ONNX的内部结构，看看一个.onnx文件里到底藏着什么秘密。

记住，ONNX不是银弹，但它绝对是现代深度学习工程化中不可或缺的一环。掌握它，你的模型就能“一次训练，到处运行”。

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