ONNX

ONNX (Open Neural Network Exchange) 是由微软和 Facebook 于 2017 年联合推出的开放神经网络互操作格式。它定义了一种标准化的计算图表示，让 PyTorch / TensorFlow / scikit-learn 训练的模型可以导出为 .onnx，然后在任何支持 ONNX 的推理引擎上运行。

核心价值：

• 跨框架互操作：训练用 PyTorch，推理用 ONNX Runtime，告别框架锁定
• 推理加速：ONNX Runtime 内置图优化（常量折叠/算子融合），速度提升 2-5x
• 多硬件部署：CPU(x86/ARM) / GPU(CUDA/TensorRT) / NPU / FPGA 统一格式
• 量化支持：INT8 / FP16 量化开箱即用，边缘设备友好
• 生态丰富：Windows ML、CoreML、TensorRT 都支持 ONNX 导入

适用场景： 模型部署、跨架构迁移、推理服务优化、边缘 AI。

环境准备

• Python：>= 3.8（推荐 3.10）
• 训练框架：PyTorch / TensorFlow 等（导出源）
• 硬件：CPU 即可运行，GPU 推理需 CUDA

安装命令

基础安装

pip install onnx onnxruntime

GPU 推理加速

# NVIDIA GPU
pip install onnxruntime-gpu

# DirectML（Windows，AMD/Intel GPU）
pip install onnxruntime-directml

模型导出依赖

# PyTorch → ONNX
pip install torch torchvision

# TensorFlow / Keras → ONNX
pip install tf2onnx

# scikit-learn → ONNX
pip install skl2onnx

验证安装

import onnx
import onnxruntime as ort

print(f"ONNX: {onnx.__version__}")
print(f"ONNX Runtime: {ort.__version__}")
print(f"可用 Provider: {ort.get_available_providers()}")

常见安装问题

Q1: onnxruntime-gpu 加载失败

CUDA 版本不匹配。检查 nvidia-smi 和 nvcc --version。降级 onnxruntime-gpu 或更新 CUDA。

Q2: 导出时报 Unsupported ONNX opset version

更新 onnx 包：pip install -U onnx。opset 版本需与 ONNX Runtime 版本匹配。

Q3: Apple Silicon 上 onnxruntime-gpu 不可用

目前仅支持 CPU provider，可使用 onnxruntime（CPU 版）或 coremltools 转 CoreML。

ONNX 完整流程：PyTorch 导出 → ONNX Runtime 推理

目标

演示标准工作流：PyTorch 训一个 ResNet → 导出为 ONNX → ONNX Runtime 推理 → 验证一致性 + 性能对比。

完整代码

import torch
import torchvision.models as models
import numpy as np
import onnx
import onnxruntime as ort
import time

# ─── 1. PyTorch 模型 ───
model = models.resnet18(pretrained=True)
model.eval()

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

# PyTorch 推理一次
with torch.no_grad():
    torch_out = model(dummy_input)

# ─── 2. 导出 ONNX ───
onnx_path = "resnet18.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,          # 保存权重
    opset_version=17,            # ONNX opset
    input_names=["input"],       # 输入名
    output_names=["output"],     # 输出名
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},  # 动态 batch
)

# 验证 ONNX 模型有效性
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print(f"✅ ONNX 模型导出成功: {onnx_path} ({onnx_model.ByteSize() / 1e6:.1f} MB)")

# ─── 3. ONNX Runtime 推理 ───
session = ort.InferenceSession(
    onnx_path,
    providers=["CPUExecutionProvider"],  # 或 ["CUDAExecutionProvider"]
)

# 准备输入
ort_inputs = {session.get_inputs()[0].name: dummy_input.numpy()}

# 预热 + 推理
_ = session.run(None, ort_inputs)
start = time.time()
for _ in range(100):
    session.run(None, ort_inputs)
ort_time = time.time() - start

# ─── 4. 验证一致性 ───
ort_out = session.run(None, ort_inputs)[0]
diff = np.abs(torch_out.numpy() - ort_out).max()
print(f"PyTorch vs ONNX 最大误差: {diff:.2e}")

# ─── 5. 性能对比 ───
with torch.no_grad():
    _ = model(dummy_input)
    start = time.time()
    for _ in range(100):
        model(dummy_input)
    torch_time = time.time() - start

print(f"\n性能对比 (100 次推理):")
print(f"  PyTorch:      {torch_time:.3f}s")
print(f"  ONNX Runtime: {ort_time:.3f}s")
print(f"  加速比:       {torch_time/ort_time:.2f}x")

运行步骤

pip install torch torchvision onnx onnxruntime numpy
python onnx_export_infer.py

预期输出

✅ ONNX 模型导出成功: resnet18.onnx (44.7 MB)
PyTorch vs ONNX 最大误差: 1.19e-07

性能对比 (100 次推理):
  PyTorch:      2.834s
  ONNX Runtime: 1.521s
  加速比:       1.86x

ONNX 入门教程：图优化、量化与部署

1. ONNX 不是另一个推理框架

ONNX 是一个格式标准——它描述"计算图长什么样"。真正的推理由 Runtime 执行：

训练框架              ONNX 格式              部署目标
─────────           ──────────            ───────────
PyTorch ────┐                             ONNX Runtime (CPU/GPU)
TensorFlow ─┼──→  model.onnx  ──→  TensorRT (NVIDIA)
scikit-learn┘                         OpenVINO (Intel)
                                       CoreML (Apple)
                                       WebNN (浏览器)

2. 图优化：ONNX Runtime 的加速原理

session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession("model.onnx", session_options)

三个优化级别：

级别	典型优化
Basic	常量折叠、冗余节点消除
Extended	算子融合（Conv+BN+ReLU → 单算子）
Layout	NCHW → NHWC 转换（适配特定硬件）

3. 量化：FP32 → INT8

from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    "resnet18.onnx",           # 输入
    "resnet18_int8.onnx",      # 输出（~1/4 大小）
    weight_type=QuantType.QInt8,
)
print("量化完成！模型缩小 75%")

量化类型	精度损失	大小	速度
FP32（原始）	0	100%	1x
FP16	<0.1%	50%	1.5-2x
INT8 动态	<0.5%	25%	1.5-2x
INT8 静态	<1%	25%	2-4x

4. 多平台部署速查

# TensorRT 加速 (NVIDIA)
trtexec --onnx=model.onnx --saveEngine=model.trt --fp16

# OpenVINO (Intel)
mo --input_model model.onnx --output_dir openvino_model/

# CoreML (Apple)
coremltools.converters.onnx.convert(model="model.onnx")

# Web (浏览器)
# 使用 onnxruntime-web + WebAssembly

5. 调试技巧

import onnx

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

# 查看算子列表
op_types = set()
for node in model.graph.node:
    op_types.add(node.op_type)
print(f"算子类型: {op_types}")

# 查看输入输出 shape
for inp in model.graph.input:
    shape = [d.dim_value for d in inp.type.tensor_type.shape.dim]
    print(f"输入: {inp.name} {shape}")

6. 常见坑

问题	原因	解决
opset 不兼容	导出用的 opset 版本 Runtime 不支持	降低 opset 或升级 Runtime
动态 shape 失败	某些后端不支持动态轴	导出时固定 batch size
自定义算子无法导出	PyTorch 中使用非标准操作	用 torch.onnx.is_in_onnx_export() 分支处理

思考题

1. ONNX 的图优化与 PyTorch 的 torch.jit.trace 有何本质区别？
2. 为什么 INT8 量化造成的精度损失通常小于 1%？原理是什么？
3. ONNX 和 TorchScript 都是"模型序列化格式"，什么时候选 ONNX？