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 更快。
常见故障定位顺序:
- 加载失败:检查文件完整性、external data、ONNX/IR/opset。
- 输入名或 shape 错误:读取 session.get_inputs(),不要猜名字。
- 算子不支持:定位节点、domain、opset 和目标 EP。
- 输出错误:先检查预处理、布局、dtype、推理模式和 tokenizer,再看导出图。
- 误差超阈值:保存同一批输入,比较中间层或分阶段输出。
- GPU 变慢:检查 CPU fallback、host-device copy、动态 shape 重新编译或 batch 太小。
- 端侧崩溃:检查模型大小、临时张量、线程、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 当作可验证、可版本化的部署工件,并在目标硬件上重新验收。
参考文献
-
ONNX 官方文档,《Introduction to ONNX》:https://onnx.ai/onnx/intro/ ↩
-
ONNX 官方仓库,包含规范、算子和模型资源:https://github.com/onnx/onnx ↩
-
ONNX Runtime 官方文档,《Execution Providers》:https://onnxruntime.ai/docs/execution-providers/ ↩
-
PyTorch 官方文档,《torch.onnx》:https://pytorch.org/docs/stable/onnx.html ↩
-
Hugging Face Optimum 官方文档,《Export a model to ONNX with Optimum》:https://huggingface.co/docs/optimum/main/en/onnx/usage_guides/export_a_model ↩
-
tensorflow-onnx 官方仓库:https://github.com/onnx/tensorflow-onnx ↩
-
Hugging Face,《Safetensors》:https://github.com/huggingface/safetensors ↩
-
ggml-org/llama.cpp,《GGUF format and gguf-py》:https://github.com/ggml-org/llama.cpp/tree/master/gguf-py ↩
-
ONNX Runtime 官方文档,《Deploy on mobile》:https://onnxruntime.ai/docs/tutorials/mobile/ ↩
-
ONNX Runtime 官方文档,《ONNX Runtime Web》:https://onnxruntime.ai/docs/get-started/with-javascript.html ↩
-
ONNX Runtime 官方文档,《Quantize ONNX models》:https://onnxruntime.ai/docs/performance/model-optimizations/quantization.html ↩
-
ONNX Runtime 官方文档,《Python》:https://onnxruntime.ai/docs/get-started/with-python.html ↩