What is ONNX（开放神经网络交换格式）?

Definition

ONNX（Open Neural Network Exchange，开放神经网络交换）是一种用于将机器学习模型表示为计算图的开源格式。最初由 Microsoft 和 Facebook（现 Meta）开发，ONNX 定义了一组通用的算子和数据类型，允许模型从一个框架（如 PyTorch）导出并导入到另一个框架（如 TensorFlow 或 ONNX Runtime 等专用推理引擎）。该格式将模型表示为有向无环图，其中节点是操作（卷积、矩阵乘法、注意力机制），边是在它们之间流动的张量。

ONNX 充当多样化 ML 生态系统之间的桥梁。不同团队可能偏好不同的训练框架，但部署环境通常有严格的要求——移动应用可能需要 Core ML，嵌入式设备可能需要 TensorRT，云服务可能使用 ONNX Runtime。ONNX 使单个训练好的模型能够通过格式转换部署到所有这些环境，无需重新训练或手动重新实现。

对于语言模型，ONNX 在推理优化方面尤为重要。ONNX Runtime（Microsoft 的推理引擎）应用图级优化——算子融合、常量折叠、内存规划——可以显著加速推理。这些优化应用于 ONNX 图表示，使其与框架无关。ONNX 生态系统还支持量化工具，可以直接在 ONNX 图上将模型压缩到 INT8 或 INT4 精度。

Why It Matters

在企业环境中，训练框架和部署环境通常是不同的。研究团队可能在 PyTorch 中开发模型，但生产系统可能需要 TensorRT 用于 GPU 推理、ONNX Runtime 用于 CPU 推理或 Core ML 用于 Apple 设备。没有像 ONNX 这样的可移植格式，每个部署目标都需要重新实现模型——这是一个成本高昂且容易出错的过程。

ONNX Runtime 由于其性能已成为生产部署的关键推理引擎。在 CPU 工作负载上，经过优化的 ONNX Runtime 可以比原生 PyTorch 推理快 2-4 倍。这种性能优势使 ONNX 转换成为生产部署管道中的标准步骤，尤其是对于需要在严格延迟要求下大规模服务的模型。

How It Works

ONNX 转换通过追踪或脚本化模型的计算图来工作。在 PyTorch 中，torch.onnx.export() 通过模型运行一个样本输入，记录所有执行的操作，并将生成的计算图序列化为 ONNX 格式。导出的图既捕获了模型架构（操作及其连接），也捕获了训练好的权重（作为图中的常量张量）。

ONNX 图随后可以被任何兼容的运行时加载。ONNX Runtime 应用优化遍：融合相邻操作（例如将 MatMul + Add + GELU 合并为单个融合内核）、消除冗余计算，以及规划内存分配以最小化开销。优化后的图被编译为目标硬件（CPU、GPU 或专用加速器）的可执行内核，生成以最小开销处理输入的推理会话。

Example Use Case

一家公司在 PyTorch 中训练了用于客户意图检测的文本分类模型。在生产环境中，他们需要在 CPU 实例上运行模型以控制成本。他们将模型导出为 ONNX，应用 ONNX Runtime 的图优化和 INT8 量化，实现了比原生 PyTorch CPU 推理 3.5 倍的吞吐量提升——使他们能够在单个实例上每秒服务 10,000 个请求，而无需三个实例。

Key Takeaways

• ONNX 是跨 ML 框架的可移植模型表示的开放标准。
• 它支持在一个框架中训练并在另一个框架中部署，无需重新实现。
• ONNX Runtime 提供图级优化，显著加速推理。
• 该格式将模型表示为包含操作和张量的有向无环计算图。
• ONNX 对于 CPU 推理优化和跨平台部署特别有价值。