1. ONNX基础认知：什么是ONNX、为什么需要ONNX、ONNX的生态与工具链概览

大家好，我是你们的老朋友。今天咱们正式开讲《PyTorch模型转ONNX避坑指南》的第一章。

说实话，我刚开始接触ONNX的时候，也一头雾水。那时候团队里有个项目，模型在PyTorch里跑得挺好，一到部署就各种报错。后来我才发现，ONNX就是那个能帮我们打通「训练」和「部署」之间那堵墙的关键角色。

好，咱们不废话，直接进入正题。

1.1 什么是ONNX？

ONNX，全称是Open Neural Network Exchange。翻译过来就是「开放神经网络交换格式」。说白了，它就是一种模型中间表示格式。

你想想看，PyTorch有自己的模型格式，TensorFlow也有自己的。它们之间互不兼容。这就好比你有台iPhone，朋友给你发了个安卓的APK安装包——你打不开啊。

ONNX就是那个「通用语言」。不管你用什么框架训练模型，只要导出成ONNX格式，就能被各种推理引擎和硬件平台识别。

我举个例子。你在PyTorch里写了个卷积层：

import torch.nn as nn

conv = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)

导出成ONNX后，它会变成类似这样的描述：

node {
  input: "input"
  input: "weight"
  input: "bias"
  output: "output"
  op_type: "Conv"
  attribute {
    name: "kernel_shape"
    ints: 3
    ints: 3
  }
  attribute {
    name: "strides"
    ints: 1
    ints: 1
  }
  attribute {
    name: "pads"
    ints: 1
    ints: 1
    ints: 1
    ints: 1
  }
}

嗯，这里要注意。ONNX的算子定义非常严格。每个算子的输入、输出、属性都有明确规范。这也是为什么有时候PyTorch模型转ONNX会报错——因为PyTorch里有些骚操作，ONNX标准里没有对应的算子。

1.2 为什么需要ONNX？

这个问题我问过很多刚入行的朋友。他们的第一反应是：「我直接用PyTorch部署不行吗？」

行，当然行。但现实往往没那么简单。

我给大家列几个真实场景：

• 跨平台部署：你的模型可能要在Windows服务器、Linux嵌入式设备、甚至手机上跑。PyTorch的运行时在这些平台上不一定都支持得好。
• 硬件加速：NVIDIA的TensorRT、Intel的OpenVINO、ARM的NPU——这些硬件加速器都只认ONNX或者自己特定的格式。你总不能让每个硬件厂商都去适配PyTorch吧？
• 模型交换：团队里有人用PyTorch，有人用TensorFlow。大家想共享一个模型，怎么办？ONNX就是那个中间人。
• 生产环境稳定性：PyTorch版本更新频繁，今天能跑的代码，明天升级了可能就炸了。ONNX格式相对稳定，更适合生产环境。

说白了，ONNX就是深度学习领域的「通用接口」。它让模型不再被某个框架绑定，真正实现了「一次训练，到处部署」。

1.3 ONNX的生态与工具链概览

ONNX不是孤零零的一个格式。它背后有一套完整的工具链。我习惯把它分成三层：

层级	工具/库	作用
模型导出层	torch.onnx、tf2onnx、onnx.export	将训练好的模型导出为ONNX格式
模型优化层	onnx-simplifier、onnxoptimizer	简化计算图、去除冗余节点、融合算子
推理部署层	ONNX Runtime、TensorRT、OpenVINO	加载ONNX模型并进行高效推理

咱们一个一个说。

模型导出层

这是最基础的一步。PyTorch官方提供了torch.onnx.export函数。你只需要传入模型、输入张量、输出路径，就能生成一个.onnx文件。

举个例子：

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

这段代码看着简单，但坑不少。比如dynamic_axes这个参数，很多人第一次用都会忽略。结果导出后模型只能处理固定batch size的输入，部署时一换batch size就报错。

模型优化层

导出的ONNX模型往往不是最优的。里面可能有很多冗余节点，比如Cast、Reshape这些。用onnx-simplifier可以一键简化：

python -m onnxsim resnet18.onnx resnet18_simplified.onnx

我遇到过最夸张的一次，一个模型简化前有2000多个节点，简化后只剩800多个。推理速度直接翻倍。

推理部署层

这是最终目的。ONNX Runtime是微软开源的推理引擎，支持CPU、GPU、甚至手机端。你只需要加载ONNX模型，然后调用run()方法就能推理。

代码示例：

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("resnet18_simplified.onnx")
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name

input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = sess.run([output_name], {input_name: input_data})

你看，就这么几行代码。不需要PyTorch，不需要GPU，甚至可以在树莓派上跑。

1.4 小结

这一章咱们聊了ONNX是什么、为什么需要它、以及它的工具链。说白了，ONNX就是深度学习模型的「通用语言」。它让模型不再被某个框架绑架，真正实现了跨平台、跨硬件的部署。

下一章，我会带大家深入PyTorch转ONNX的细节。咱们聊聊那些常见的坑，以及怎么避开它们。

嗯，今天就到这儿。有问题欢迎留言，咱们下期见。