ONNX 是一种跨框架、跨语言、跨硬件的推理计算图交换格式。它不是新的神经网络,也不是训练框架,更不是把任意权重文件改名后的“万能格式”。训练框架负责定义模型和学习参数,ONNX 负责用统一方式描述推理图,运行时负责在具体硬件上执行或编译这张图。

本文从概念、格式、转换、部署和工程验收五个层次展开,重点回答:ONNX 文件保存什么;PyTorch、TensorFlow、SafeTensors、GGUF 如何转 ONNX;如何部署到 GPU、手机和浏览器;以及什么时候不该选择 ONNX。文中的代码以“可理解、可检查”为优先,真实项目中要替换成自己的输入、opset、精度阈值和性能目标。

一、先建立正确的模型文件观

1. 训练框架、权重文件、交换格式、运行时不是一回事

概念 解决的问题 常见例子
训练/建模框架 定义网络、求梯度、训练和保存 PyTorch、TensorFlow、Keras
权重序列化 把参数张量写入磁盘 .pth、.pt、Checkpoint、SafeTensors
模型交换格式 描述计算图、算子、输入输出和权重 ONNX
推理运行时/编译器 在具体硬件执行图 ONNX Runtime、TensorRT、OpenVINO、Core ML

.pth 或 .safetensors 可能只包含参数;而 ONNX 通常同时包含“怎么计算”的图和“计算所需”的权重。ONNX Runtime 是执行 ONNX 的程序,不是 ONNX 文件本身。

可以把它类比成:SafeTensors 是一箱标好名字和形状的零件;ONNX 是一张装配图,里面还可以内嵌这些零件;ONNX Runtime 是按照装配图运行的机器。仅仅拥有参数张量,并不自动拥有网络结构。

2. ONNX 是什么

ONNX(Open Neural Network Exchange)是开放的机器学习模型表示格式和生态项目。它用相对独立于训练框架的方式表达输入输出、张量类型和形状、节点及算子类型、数据依赖、常量和权重,以及版本和元数据。1

从抽象上看,ONNX 是一张有向图:

输入张量 x
    |
    v
[Conv] -> [BatchNormalization] -> [Relu] -> [GlobalAveragePool]
    |                                             |
    +----------------------> [Gemm] ------------> 输出 y

节点是计算,边是张量。模型文件描述的是一次前向推理,而不是训练数据、损失函数、优化器状态和反向传播过程。

3. 文件内部大致长什么样

ONNX 主文件通常是 Protocol Buffers(protobuf)序列化的 ModelProto,应使用 onnx Python 包检查。2

ModelProto
├── ir_version                 # ONNX 内部表示版本
├── opset_import               # 标准算子集版本
├── graph
│   ├── input / output         # 图的接口
│   ├── node                   # Conv、MatMul、Add 等节点
│   ├── initializer            # 权重或常量
│   └── value_info             # 中间张量信息
└── metadata_props             # 可选元数据

权重小的时候放在 initializer;大权重可以作为 external data 外置。外置数据必须与主文件一起发布。

4. IR、opset 与 Execution Provider

  • IR version:ONNX 文件内部结构的版本。
  • Opset version:标准算子的版本。同名算子在不同 opset 中可能有不同属性或语义。
  • Execution Provider(EP):ONNX Runtime 执行后端,例如 CPU、CUDA、TensorRT、DirectML、CoreML、OpenVINO、NNAPI、NPU 等。3

onnx.checker 检查通过,只说明模型结构符合规范,不代表目标 EP 支持所有节点,也不代表性能一定更好。

二、为什么需要 ONNX

1. 把训练与推理解耦

训练环境需要 Python、自动求导、数据加载、分布式训练等复杂依赖;线上推理更关心启动、内存、吞吐、延迟和稳定升级。导出 ONNX 后,服务端可以只部署推理运行时,也方便 C++、C#、Java、Rust 等语言调用同一模型。

2. 连接不同硬件后端

PyTorch / TensorFlow / Keras
            |
            v
          ONNX
       /    |       \
CPU + ORT  CUDA   TensorRT
       |      |        |
服务器   NVIDIA GPU  专用高性能引擎

这不保证“一次导出、处处最快”。不同硬件的最终格式可能是 TensorRT engine、OpenVINO IR、Core ML、TFLite 或厂商编译产物。ONNX 是重要的中间站和互操作边界。

3. ONNX 的边界

ONNX 不是训练 checkpoint,也不完整包含 tokenizer、词表、采样器、对话模板。动态控制流、Python 副作用和任意自定义算子可能难以表达。大语言模型的 KV cache、量化布局、采样和多轮状态需要额外设计。导出成功也不代表结果正确,更不代表一定比原框架快。

三、导出前必须确认的四件事

1. 明确推理接口

先写清输入名称、数据类型、静态或动态 shape、输出定义、预处理与后处理边界,以及目标硬件、批大小、延迟和精度目标。

图像分类模型的输入往往不是 RGB 图片,而是 resize、裁剪、归一化、HWC 到 CHW 后的 float32[N, 3, H, W] 张量。预处理不一致,输出再“正常”也没有意义。

2. 关闭训练行为

PyTorch 导出前通常需要:

import torch

model.eval()

with torch.no_grad():
    output = model(example_input)

这会关闭 Dropout 等训练期随机行为,并让 BatchNorm 使用已学习的推理统计量。

3. 固定或声明动态维度

静态 shape 容易优化;动态 shape 更灵活但需要运行时和 EP 支持。不要无脑把所有维度都设成动态。

# 只让 batch 变化;图像高宽是否动态要看目标后端。
dynamic_axes = {
    "images": {0: "batch"},
    "logits": {0: "batch"},
}

Transformer 常见动态维度是 batch 和 sequence length;带 KV cache 的解码模型还要设计 cache 的输入输出接口。

4. 选择 opset 并记录版本

opset 太旧可能缺算子语义,太新又可能超出目标运行时支持。应记录训练框架、导出脚本、ONNX/opset、示例输入、预处理、运行时、精度结果和性能基线。

四、PyTorch 转 ONNX

1. 使用 torch.onnx.export

现代 PyTorch 导出器使用 torch.export,并可通过 dynamo=True 生成 ExportedProgram。不同版本能力可能变化,因此工程脚本要显式指定。4

import torch

# 用训练完成的模型替换 MyModel。
model = MyModel().eval()
example_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    (example_input,),
    "model.onnx",
    input_names=["images"],
    output_names=["logits"],
    opset_version=18,
    dynamo=True,
    # 只让 batch 动态,为编译器保留优化空间。
    dynamic_axes={
        "images": {0: "batch"},
        "logits": {0: "batch"},
    },
)

导出失败时,先看失败节点和错误类型:可能是 Python 控制流、没有 ONNX 映射的算子,或输入样例没有覆盖真实路径。旧项目的 dynamo=False TorchScript 路径可能仍可用,但要明确记录并按目标版本测试。

2. Hugging Face Transformers 模型

Transformers 模型通常要同时处理 tokenizer、模型包装、输出对象和动态输入。Optimum 更适合这类导出:5

pip install "optimum[onnxruntime]" transformers onnx onnxruntime
optimum-cli export onnx   --model bert-base-chinese   --task text-classification   onnx/bert-base-chinese

task 必须和真实任务一致,例如 feature-extraction、sequence-classification、text-generation。生成式模型可能拆成 encoder、decoder、带 past key/value 的 decoder 等多个文件。

3. 导出后的验证

import numpy as np
import onnxruntime as ort

session = ort.InferenceSession(
    "model.onnx",
    providers=["CPUExecutionProvider"],
)
input_name = session.get_inputs()[0].name
onnx_output = session.run(None, {input_name: example_input.numpy()})[0]

with torch.inference_mode():
    torch_output = model(example_input).detach().cpu().numpy()

max_error = np.max(np.abs(torch_output - onnx_output))
print("最大绝对误差:", max_error)

分类任务还应比较 top-1/top-k;检测、分割和生成任务应比较业务输出,而不只是逐元素误差。

五、TensorFlow / Keras 转 ONNX

TensorFlow 生态常用 tf2onnx,支持 SavedModel、Keras、ConcreteFunction 等入口。6

1. 从 SavedModel 转换

pip install tf2onnx onnx onnxruntime tensorflow
python -m tf2onnx.convert   --saved-model ./saved_model   --output ./onnx/model.onnx   --opset 18

2. 从 Keras 对象转换

import tensorflow as tf
import tf2onnx

# model 是加载并准备好的 Keras 模型。
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=(
        tf.TensorSpec(
            shape=[None, 224, 224, 3],
            dtype=tf.float32,
            name="images",
        ),
    ),
    opset=18,
    output_path="model.onnx",
)

TensorFlow 常见 NHWC,PyTorch 常见 NCHW。转换器可能插入 Transpose,但调用方预处理仍必须遵循实际输入约定。TensorList、RaggedTensor、控制流或自定义 op 要逐项确认目标运行时支持。

TFLite 是另一种端侧格式。通常应从原始 TensorFlow/Keras/SavedModel 重新导出 ONNX,而不是把任意 TFLite 反向还原为完整 ONNX。TFLite 到 ONNX 的工具必须逐模型验证算子、量化和自定义实现。

六、SafeTensors 转 ONNX:先问“图在哪里”

SafeTensors 是安全、快速、可按张量读取的存储格式,主要保存张量名称、dtype、shape、偏移、原始字节和少量 metadata。它通常不保存可执行网络结构7

因此不存在可靠的通用“任意 SafeTensors 一键转 ONNX”命令。正确问题是:是否还拥有解释这些权重的模型类、配置和 tokenizer?

如果是 Hugging Face 模型,目录通常包含:

config.json
model.safetensors              # 也可能是多个分片
tokenizer.json / tokenizer.model
tokenizer_config.json

这时直接用 Optimum:

pip install "optimum[onnxruntime]" transformers onnx onnxruntime
optimum-cli export onnx   --model ./hf-model   --task text-classification   ./onnx-model

如果只有 SafeTensors 没有 config 和模型实现,需要先找到架构来源,或用完全匹配的结构实例化并加载 state dict。严格检查 missing keys、unexpected keys、dtype 和 shape;“能加载”不等于语义匹配。

SafeTensors 的安全优势不等于 ONNX 的安全保证。模型文件应限制来源、校验 hash、在隔离环境检查,并审查 custom op。

七、GGUF 转 ONNX:通常不是直接转换问题

GGUF 是 ggml/llama.cpp 生态常用的模型文件格式,把元数据、张量和量化信息组织到一起,服务于高效本地推理。它与 ONNX 的目标不同。8

  • GGUF 偏向特定推理生态中的可加载权重与量化布局;
  • ONNX 偏向通用计算图交换和多运行时执行;
  • GGUF 依赖架构元数据、实现约定和量化 kernel;
  • ONNX 需要显式表达目标运行时理解的算子图与接口。

所以通常没有“GGUF 直接转 ONNX 且保持所有量化语义”的通用可靠路径。

推荐回到原始来源:

原始 Hugging Face / PyTorch 权重 + config
    -> Transformers / PyTorch 模型
    -> Optimum 或 torch.onnx.export
    -> ONNX
    -> ONNX Runtime / TensorRT

如果只剩 GGUF,需要确认架构和版本,读取元数据与张量,把量化张量恢复或映射为目标模型格式,用匹配结构重建计算图,再导出并对齐结果。这已是针对具体架构和量化类型的移植项目,容易漏掉 RoPE、RMSNorm、权重转置、分组查询注意力、KV cache、scale/zero point 和 logits 处理。

自回归生成也不是一次 input -> output:

token_ids + attention_mask + past_key_values
             |
             v
      logits + next_key_values
             |
             v
      采样/选 token,再进入下一轮

ONNX 能表达前向计算,但 tokenizer、停止词、采样、beam search、KV cache 生命周期仍需业务或生成运行时配合。

八、ONNX 模型部署到 GPU

1. 安装和查看后端

NVIDIA CUDA 场景通常安装:

pip install onnx onnxruntime-gpu

查看运行时注册的后端:

import onnxruntime as ort

print(ort.get_available_providers())

CUDAExecutionProvider 能否真正工作,还取决于驱动、CUDA、cuDNN 与 onnxruntime-gpu 构建的兼容矩阵。包安装成功不代表 GPU 已被使用。

2. CUDA Execution Provider

import onnxruntime as ort

session = ort.InferenceSession(
    "model.onnx",
    providers=[
        ("CUDAExecutionProvider", {
            "device_id": 0,
            # 参数按模型和 ORT 版本选择,不要盲抄。
            "arena_extend_strategy": "kNextPowerOfTwo",
        }),
        "CPUExecutionProvider",
    ],
)

print(session.get_providers())

CPU 放在列表末尾表示回退,但不支持的节点可能产生跨设备拷贝和延迟抖动。生产环境应检查 profiling、日志和节点分配,而不是只看 session.run 成功。

3. TensorRT Execution Provider

providers = [
    ("TensorrtExecutionProvider", {
        "device_id": 0,
        "trt_engine_cache_enable": True,
        "trt_engine_cache_path": "./trt-cache",
    }),
    ("CUDAExecutionProvider", {"device_id": 0}),
    "CPUExecutionProvider",
]

session = ort.InferenceSession("model.onnx", providers=providers)

TensorRT 对固定 shape、静态 profile、FP16/INT8 校准和算子支持敏感。首次构建 engine 可能较慢,engine 通常和 GPU 架构、TensorRT 版本绑定。

4. GPU 部署验证

应确认可用 EP、provider 顺序、版本兼容性,并用 profiling 证明主要节点在 GPU。还要测试 CPU fallback、batch、shape、并发、warm-up、FP16/INT8 精度,以及包含数据拷贝、预处理和后处理的端到端延迟。

九、部署到端侧和浏览器

1. 移动端与嵌入式设备

ONNX Runtime 提供移动端包和 API,可在 Android、iOS 等环境调用,并在硬件支持时使用 NNAPI、Core ML 等 EP。9

端侧通常需要用 onnxruntime.tools.convert_onnx_models_to_ort 生成 ORT 格式,只保留模型实际使用的算子,并针对真实设备测试量化、内存峰值和冷启动。ORT 是针对 ONNX Runtime 的优化产物,不是通用交换格式;可把 ONNX 作为主产物,把 ORT 作为端侧发布产物。

2. 浏览器与 WebAssembly/WebGPU

ONNX Runtime Web 可在浏览器使用 WebAssembly、WebGPU、WebNN 等后端。要考虑模型下载和缓存、首次编译、浏览器支持、Worker、内存及算子覆盖。10

import * as ort from "onnxruntime-web";

const session = await ort.InferenceSession.create(
  "/models/model.onnx",
  { executionProviders: ["wasm"] },
);

const input = new ort.Tensor(
  "float32",
  float32Array,
  [1, 3, 224, 224],
);

const output = await session.run({ images: input });
console.log(output.logits);

如果使用 WebGPU,应确认版本、浏览器和设备支持,并与 WASM 后端实际对比。

3. 端侧量化

常见路线包括 FP32 到 FP16、FP32 到 INT8、权重-only、静态量化和动态量化。INT8 往往需要校准数据或量化感知训练。量化后必须重新做数值和业务指标验证;模型更小不等于端到端更快。11

十、ONNX 相关 Python SDK

不要把所有包都叫作 ONNX SDK,它们处于不同层:

  • onnxruntime:InferenceSession、输入输出、execution provider、线程和 profiling,是最常用的 Python 推理入口。12
  • tf2onnx:TensorFlow、Keras、TensorFlow.js、TFLite 到 ONNX。
  • torch.onnx:PyTorch 官方导出入口。
  • optimum[onnxruntime]:Hugging Face Transformers 导出和 ORT 集成。
  • skl2onnx:部分 scikit-learn 模型。
  • onnxmltools:部分传统机器学习工具。
  • onnxsim、onnxoptimizer:图简化和部分优化。
  • onnxruntime-genai:更贴近生成模型的封装,但不是通用 Transformer 导出器。
  • Netron:可视化 ONNX 图、输入输出和节点。

推荐先用 onnx checker 和 ONNX Runtime CPU 建立正确性基线,再切换目标 EP。优化前后都要做结构检查和数值回归。

import onnx
import onnxruntime as ort

model = onnx.load("model.onnx")
onnx.checker.check_model(model)

session = ort.InferenceSession(
    "model.onnx",
    providers=["CPUExecutionProvider"],
)

for value in session.get_inputs():
    print("输入:", value.name, value.shape, value.type)

for value in session.get_outputs():
    print("输出:", value.name, value.shape, value.type)

十一、完整的发布流水线

训练 checkpoint
      |
      v
加载原模型 + eval + 明确输入签名
      |
      v
导出 ONNX
      |
      +--> onnx.checker 结构检查
      |
      +--> ORT CPU 数值对齐
      |
      +--> 目标 GPU/端侧 EP 数值对齐
      |
      +--> 性能基准与 profiling
      |
      v
量化/图优化/编译
      |
      v
再次做结构、精度、性能和包体积验证
      |
      v
发布 ONNX 或目标后端产物

基准至少区分首次加载/编译、warm-up 后延迟、多 batch 吞吐、并发 p50/p95/p99、内存峰值、端到端耗时、GPU 利用率和数据传输开销。不要用一次 session.run 的时间宣称 GPU 更快。

常见故障定位顺序:

  1. 加载失败:检查文件完整性、external data、ONNX/IR/opset。
  2. 输入名或 shape 错误:读取 session.get_inputs(),不要猜名字。
  3. 算子不支持:定位节点、domain、opset 和目标 EP。
  4. 输出错误:先检查预处理、布局、dtype、推理模式和 tokenizer,再看导出图。
  5. 误差超阈值:保存同一批输入,比较中间层或分阶段输出。
  6. GPU 变慢:检查 CPU fallback、host-device copy、动态 shape 重新编译或 batch 太小。
  7. 端侧崩溃:检查模型大小、临时张量、线程、arena、量化和设备专属 EP。

十二、选择指南

场景 建议
PyTorch 训练,Python 服务,CPU/GPU 推理 先建立 ORT 基线,再与原生 PyTorch 对比
需要 C++/Java/移动端/浏览器 ONNX 通常值得作为交换格式
NVIDIA 固定 GPU、极低延迟 ONNX -> TensorRT 可能更合适,保留 ONNX
TensorFlow/Keras 迁移 使用 tf2onnx,验证算子和布局
只有 SafeTensors 找回 config 和模型类后导出
只有 GGUF 且目标是 llama.cpp 继续使用 GGUF,不必为了统一转 ONNX
只有 GGUF 但目标是 TensorRT/ORT 尽量找回原始权重,否则按具体架构移植
需要完整训练续跑 保留原 checkpoint,ONNX 只作为推理产物
复杂控制流或大量 custom op 先评估目标运行时,必要时保留原框架

最重要的结论是:ONNX 解决的是推理图的互操作和部署问题,不是所有模型文件之间的万能翻译。 可行性取决于原文件是否包含结构、是否有导出器、目标 opset/EP 是否支持,以及预处理和状态接口是否完整。工程上应保留原始权重和配置,把 ONNX 当作可验证、可版本化的部署工件,并在目标硬件上重新验收。

参考文献

  1. ONNX 官方文档,《Introduction to ONNX》:https://onnx.ai/onnx/intro/ 

  2. ONNX 官方仓库,包含规范、算子和模型资源:https://github.com/onnx/onnx 

  3. ONNX Runtime 官方文档,《Execution Providers》:https://onnxruntime.ai/docs/execution-providers/ 

  4. PyTorch 官方文档,《torch.onnx》:https://pytorch.org/docs/stable/onnx.html 

  5. Hugging Face Optimum 官方文档,《Export a model to ONNX with Optimum》:https://huggingface.co/docs/optimum/main/en/onnx/usage_guides/export_a_model 

  6. tensorflow-onnx 官方仓库:https://github.com/onnx/tensorflow-onnx 

  7. Hugging Face,《Safetensors》:https://github.com/huggingface/safetensors 

  8. ggml-org/llama.cpp,《GGUF format and gguf-py》:https://github.com/ggml-org/llama.cpp/tree/master/gguf-py 

  9. ONNX Runtime 官方文档,《Deploy on mobile》:https://onnxruntime.ai/docs/tutorials/mobile/ 

  10. ONNX Runtime 官方文档,《ONNX Runtime Web》:https://onnxruntime.ai/docs/get-started/with-javascript.html 

  11. ONNX Runtime 官方文档,《Quantize ONNX models》:https://onnxruntime.ai/docs/performance/model-optimizations/quantization.html 

  12. ONNX Runtime 官方文档,《Python》:https://onnxruntime.ai/docs/get-started/with-python.html