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

1.1 什么是ONNX？说白了就是个模型“通用语言”

ONNX，全称是Open Neural Network Exchange。名字挺长，但核心就一句话：它是深度学习模型的“通用交换格式”。

我刚开始接触AI部署那会儿，最头疼的就是模型格式不统一。PyTorch训练好的模型，想部署到TensorRT上？得先转成.pt，再转成.wts，再转成.engine……中间但凡有个版本对不上，直接报错。后来ONNX出现了，大家约定好：你训练完，先导出成.onnx，剩下的部署工作我来搞定。

说白了，ONNX就像USB接口。不管你的设备是苹果、安卓还是Windows，只要插上USB，就能传输数据。ONNX就是AI模型界的USB。

核心定义：ONNX是一种开放的、跨平台的深度学习模型表示格式。它定义了一套标准化的算子集和计算图描述方式，让模型可以在不同框架之间自由迁移。

1.2 为什么需要ONNX？——我踩过的坑告诉你

你可能觉得：“我用PyTorch训练，用PyTorch推理不就行了？干嘛多此一举？”

嗯，这个问题我当年也问过。直到我在项目中遇到这几个场景：

  • 场景一：训练用PyTorch，部署用TensorRT——PyTorch的推理速度在GPU上还行，但到了边缘设备（比如Jetson Nano）上，TensorRT能快3-5倍。没有ONNX，你得手动写转换脚本，一个算子一个算子地调。
  • 场景二：团队里有人用PyTorch，有人用TensorFlow——模型交接时，对方说“你转成.pb给我”，你心里一万个不情愿。有了ONNX，大家统一导出.onnx，谁都不用迁就谁。
  • 场景三：模型要跑在手机端（CoreML）或浏览器端（WebGL）——这些平台只认自己的格式。ONNX就像中间人，帮你把模型“翻译”成目标平台能懂的语言。

我曾经有个项目，客户要求模型必须跑在华为昇腾芯片上。昇腾的推理引擎只支持ONNX和Caffe。我花了三天把PyTorch模型转成ONNX，又花了一天调精度对齐。如果没有ONNX，我可能得重写整个推理代码——那画面太美我不敢看。

我的建议：无论你用什么框架训练，养成导出ONNX的习惯。哪怕你现在用不到，将来某天客户突然说“我要部署在xxx上”，你会感谢自己当初多存了一份.onnx文件。

1.3 ONNX的生态与优势——它凭什么这么火？

ONNX能火，不是靠吹的。它的生态，说白了就是“大家都来玩”。

1.3.1 框架支持：几乎全覆盖

框架 导出ONNX 导入ONNX 备注
PyTorch ✅ torch.onnx.export ✅ onnx2pytorch 官方支持，最成熟
TensorFlow ✅ tf2onnx ✅ onnx2tf 需要额外工具
Keras ✅ keras2onnx ✅ onnx2keras 社区维护
MXNet ✅ mxnet2onnx ✅ onnx2mxnet 官方支持
PaddlePaddle ✅ paddle2onnx ✅ onnx2paddle 百度官方支持

1.3.2 推理引擎支持：你想得到的都有

ONNX模型训练好之后，能跑在哪些推理引擎上？我列几个常用的：

  • ONNX Runtime：微软出品，CPU/GPU都支持，性能优秀。我个人最常用这个做快速验证。
  • TensorRT：NVIDIA的GPU加速引擎。ONNX转TensorRT，是很多高性能部署场景的标准流程。
  • OpenVINO：Intel的推理引擎，专门优化CPU和集成显卡。我在边缘设备上经常用。
  • CoreML：苹果的移动端推理引擎。ONNX转CoreML，就能跑在iPhone上。
  • NCNN / MNN：腾讯和阿里巴巴的移动端引擎。ONNX转NCNN，是很多Android端AI应用的标配。

一句话总结：ONNX的生态，就是“一次导出，到处运行”。你只需要维护一份模型文件，剩下的交给ONNX生态去适配。

1.3.3 ONNX的优势——为什么我推荐你学？

除了生态，ONNX本身还有几个硬核优势：

  1. 标准化算子集：ONNX定义了100多个标准算子（Conv、Relu、BatchNorm等）。不同框架的算子，只要映射到ONNX标准，就能互通。我遇到过有些框架的算子实现有细微差异，ONNX帮我们统一了行为。
  2. 计算图优化：ONNX模型可以做一些图级别的优化，比如常量折叠、算子融合。这些优化在推理时能明显提速。我记得有一次，一个ResNet50模型，ONNX Runtime比PyTorch原生推理快了15%。
  3. 量化支持：ONNX支持INT8、FP16等低精度量化。对于边缘设备，量化后模型大小能缩小4倍，速度提升2-3倍。我做过一个项目，模型从FP32量化到INT8，精度只掉了0.3%，但推理速度翻了一倍。
  4. 跨平台：Windows、Linux、macOS、Android、iOS……ONNX都能跑。你想想看，一个模型文件，从服务器到手机，从云端到边缘，全都能用，这得多省事。

注意：ONNX不是万能的。有些自定义算子（比如你自己写的CUDA kernel）可能不支持。遇到这种情况，要么用ONNX的自定义算子扩展，要么老老实实写原生推理代码。我建议：训练时尽量用标准算子，少搞花活。

1.4 一个简单的ONNX导出示例

光说不练假把式。我写个最简单的PyTorch模型导出ONNX的例子，你感受一下：

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(16, 10)

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

# 创建模型和虚拟输入
model = SimpleModel()
dummy_input = torch.randn(1, 3, 224, 224)

# 导出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模型导出成功！")

这段代码跑完，你会在当前目录下看到一个simple_model.onnx文件。用Netron打开它，就能看到完整的计算图结构。嗯，是不是很简单？

小技巧：导出时加上dynamic_axes参数，可以让模型支持动态batch size。这样部署时，你就不用固定输入尺寸了。我习惯在导出时都加上这个，省得后面改。

1.5 本章小结

这一章我们聊了：

  • ONNX是什么——模型界的“通用语言”
  • 为什么需要ONNX——解决框架碎片化、跨平台部署的痛点
  • ONNX的生态——框架支持、推理引擎支持、标准化优势
  • 一个简单的导出示例——让你亲手感受ONNX的便利

下一章，我会带你深入ONNX的内部结构，看看一个.onnx文件里到底藏着什么。到时候，你会对“计算图”这个概念有更直观的理解。

记住：ONNX不是终点，它是连接训练和部署的桥梁。把这座桥搭好了，后面的路就好走了。

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