ONNX：在任意框架训练，在任意硬件部署

您用 PyTorch 训练了一个模型，但生产系统运行的是 TensorFlow。需要从零开始重写一切吗？

不，完全不必。

模型互操作性常常让机器学习工程师头疼。不同框架使用不同格式，互转时可能会破坏模型或引入意外行为。您既不想要这些问题，也不想维护同一模型的多个版本。

ONNX（Open Neural Network Exchange）通过为机器学习模型提供通用格式来解决这一问题。它让您可以在一个框架中训练，在另一个框架中部署，而无需重写代码或处理错误。

本文将向您展示如何将模型转换为 ONNX 格式，如何使用 ONNX Runtime 进行推理，如何为生产优化模型，以及如何在从边缘设备到云服务器的各种平台上部署。

ONNX 入门

使用 ONNX 前，您需要先配置开发环境。

本节将带您完成所需的一切——从软件要求到创建干净、可复现的工作区。我会演示如何在 Windows、Linux 和 macOS 上安装 ONNX 和 ONNX Runtime。

先决条件

ONNX 可在所有主流操作系统上使用。

您的机器上需要安装 Python 3.8 或更高版本。基础要求就这些。ONNX 支持 Python 的 ABI3（应用二进制接口），这意味着您可以直接使用预编译的二进制包，无需从源码编译。

您将安装以下内容：

uv（在 astral.sh/uv 获取）
Python 3.8+（uv 会为您处理）
pip（随 uv 一起提供）

想进一步了解 uv 的工作原理以及为何它值得作为首选，请阅读我们的最快的 Python 包管理器指南。

Windows 用户可使用 PowerShell 安装 uv：

powershell -c "irm <https://astral.sh/uv/install.ps1> | iex"技术支持

如果遇到执行策略错误，请先以管理员身份运行 PowerShell。

Linux 用户可通过以下命令安装 uv：

curl -LsSf <https://astral.sh/uv/install.sh> | sh技术支持

适用于任意 Linux 发行版。

最后，macOS 用户可通过 Homebrew 安装 uv，或使用 curl 命令：

brew install uv

# or

curl -LsSf <https://astral.sh/uv/install.sh> | sh技术支持

环境配置

使用 uv 创建项目的方法如下：

mkdir onnx-project
cd onnx-project

# Initialize a uv project with Python 3.10
uv init --python 3.13技术支持

uv 会创建一个 .venv 目录和一个 pyproject.toml 文件。通过 uv 运行命令时，虚拟环境会自动激活。

现在安装 ONNX 和 ONNX Runtime：

uv add onnx onnxruntime

# Verify the installation
uv run python -c "import onnx; print(onnx.__version__)"
uv run python -c "import onnxruntime; print(onnxruntime.__version__)"技术支持

验证命令应会打印版本号：

您的依赖已固定。uv 会自动将所有包版本写入 uv.lock。这使构建可复现——任何人只需运行 uv sync 就能重建与您完全一致的环境。

以下是 pyproject.toml 的样子：

[project]
name = "onnx-project"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
    "onnx>=1.19.1",
    "onnxruntime>=1.23.2",
]技术支持

就这样，环境已配置完成！接下来进入基础知识。

理解 ONNX 基础

开始转换模型前，您需要先了解 ONNX 的工作方式。

本节将拆解 ONNX 格式并解释其基于图的架构。我会展示 ONNX 文件内部结构以及数据在其中的流动方式。

ONNX 模型格式

ONNX 模型是一个包含神经网络结构和权重的文件。

可将其视为蓝图。文件描述要执行哪些操作、以何种顺序、使用哪些参数。训练好的权重与这份蓝图一同存储，因此模型无需额外文件即可运行。

ONNX 于 2017 年由微软和 Facebook（现 Meta）合作推出。目标很简单——每次切换框架时别再重写模型。尽管 PyTorch、TensorFlow 等框架不断进步，但模型在它们之间迁移离不开手动转换。

ONNX 改变了这一局面。1.0 版本支持计算机视觉和简单神经网络的基础算子。如今的 ONNX 已支持 Transformer、大语言模型，以及 2017 年时尚不存在的复杂架构。

它使用 Protocol Buffers（protobuf）来序列化数据。Protobuf 是 Google 为高效存储与传输数据而开发的二进制格式。读写速度远快于 JSON 或 XML，文件也更小。两全其美。

一个 ONNX 模型包含以下内容：

图：网络结构
权重：训练得到的模型参数
元数据：版本信息、生产者名称和模型文档
Opset 版本：可用的操作集合

Opset 版本很重要。ONNX 会持续演进，不断加入新操作并改进现有操作。您的模型文件会声明所用的 opset 版本，以便运行时知道如何执行。

关键概念与架构

ONNX 以计算图的形式表示您的模型。

图由节点和边组成。节点是操作（如矩阵乘法或激活函数），边是张量（您的数据），在这些操作之间流动。ONNX 正是通过这种方式来描述“先相乘这些矩阵，再应用 ReLU，然后再相乘”。

来看一个简单示例：

import onnx
from onnx import helper, TensorProto

# Create input and output tensors
input_tensor = helper.make_tensor_value_info(
    "input", TensorProto.FLOAT, [1, 3, 224, 224]
)
output_tensor = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1000])

# Create a node (operation)
node = helper.make_node(
    "Relu",  # Operation type
    ["input"],  # Input edges 
    ["output"],  # Output edges
)

# Create the graph
graph = helper.make_graph(
    [node],  # List of nodes
    "simple_model",  # Graph name
    [input_tensor],  # Inputs
    [output_tensor],  # Outputs
)

# Create the model
model = helper.make_model(graph)
onnx.save(model, "simple_model.onnx")技术支持

每个节点执行一个操作。操作类型（如 Relu、Conv 或 MatMul）来源于 ONNX 算子集合。您不能随意编造操作名称——它们必须存在于所用 opset 版本中。

边连接节点并承载张量。张量具有形状和数据类型。当您定义 [1, 3, 224, 224] 时，表示“该张量有 4 个维度，大小分别为这些数”。运行时会使用这些信息分配内存并校验计算图。

图是单向流动的。数据从输入节点进入，经过各个操作，最终从输出节点流出。不允许存在环路——ONNX 不直接支持循环连接。您需要展开循环或使用专门面向序列的操作。

这种图结构实现了框架间互操作。PyTorch 倾向于即时执行，TensorFlow 使用静态图，scikit-learn 则是完全不同的抽象。但它们都能导出为相同的图格式。

图也使共享优化成为可能。ONNX 优化器可以进行算子融合（将多个节点合并为一个）、消除死代码（移除未使用的节点），或用更快的等价操作替换慢操作。这些优化不依赖于创建模型的具体框架。

接下来，我将演示如何将模型转换为 ONNX 格式的具体流程。

将模型转换为 ONNX 格式

本节将演示如何从 PyTorch、TensorFlow 和 scikit-learn 导出模型。我会带您完成转换流程，并展示如何验证转换后的模型是否正常工作。

框架支持与工具

ONNX 支持最流行的机器学习框架。

PyTorch 通过 torch.onnx.export() 内置支持 ONNX 导出。由于 PyTorch 和 ONNX 都由 Meta 维护，这是最成熟的转换器。

TensorFlow 使用 tf2onnx 进行转换。它是一个独立库，可将 TensorFlow 计算图转换为 ONNX 格式。通过 uv add tf2onnx 安装即可使用。

scikit-learn 通过 skl2onnx 转换。该库处理传统机器学习模型，如随机森林、线性回归和 SVM。尽管深度学习框架更受关注，scikit-learn 模型也同样适用。

以上三者最为常见，但您也可以使用其他受支持的框架——只需通过 uv add <name>: 安装相应依赖：

Keras：使用 tf2onnx（Keras 是 TensorFlow 的一部分）
XGBoost：使用 onnxmltools
LightGBM：使用 onnxmltools
MATLAB：内置 ONNX 导出
Paddle：使用 paddle2onnx

每个框架都有自己的转换器，因为它们对模型的表示不同。PyTorch 使用动态图，TensorFlow 使用静态图，而 scikit-learn 根本不使用图。转换器会将这些表示翻译为 ONNX 的图格式。

您也可以使用他人训练的模型。 ONNX Model Zoo 托管了可直接使用的预训练模型。您可以找到计算机视觉模型（ResNet、YOLO、EfficientNet）、NLP 模型（BERT、GPT-2）等。可从官方 GitHub 仓库下载。

Hugging Face 也托管 ONNX 模型。搜索名称中包含“onnx”的模型，或按 ONNX 格式筛选。许多流行的 Transformer 模型已提供经推理优化的 ONNX 版本。

继续之前，请运行此命令，通过 uv 获取所需的全部库：

uv add torch torchvision tensorflow tf2onnx scikit-learn skl2onnx onnxscript技术支持

分步转换流程

以下是将 PyTorch 模型转换为 ONNX 的方法：

import torch
import torch.nn as nn

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

# Create and prepare the model
model = SimpleModel()
model.eval()  # Set to evaluation mode

# Create dummy input with the same shape as your real data
dummy_input = torch.randn(1, 10)

# Export to ONNX
torch.onnx.export(
    model,  # Model to export
    dummy_input,  # Example input
    "simple_model.onnx",  # Output file
    input_names=["input"],  # Name for the input
    output_names=["output"],  # Name for the output
    dynamic_shapes={"x": {0: "batch_size"}},  # Allow variable batch size
)技术支持

初识 PyTorch？别被它拦住脚步。我们的 Deep Learning with PyTorch 课程可在数小时内带您掌握基础。

模拟输入很关键。ONNX 需要知道输入形状来构建计算图。如果您的模型接收变长序列或不同批大小，请使用 dynamic_shapes 将这些维度标记为动态。

导出前务必将模型设为评估模式（model.eval()）。这会禁用 dropout 和批归一化的训练行为。若忘记此步骤，转换后的模型将与原模型不一致。

动态形状可能带来问题。如果转换器提示未知维度，请确保正确指定 dynamic_shapes，或提供具体形状。

接下来切换到 TensorFlow，您可以运行以下替代代码：

import numpy as np
# Temporary compatibility fixes for tf2onnx
if not hasattr(np, "object"):
    np.object = object
if not hasattr(np, "cast"):
    np.cast = lambda dtype: np.asarray

import tensorflow as tf
import tf2onnx

# Define a simple Keras model (same spirit as your PyTorch one)
inputs = tf.keras.Input(shape=(10,), name="input")
outputs = tf.keras.layers.Dense(5, name="output")(inputs)
model = tf.keras.Model(inputs, outputs)

# Save the model (optional, if you want a .h5 file)
model.save("simple_tf_model.h5")

# Define a dynamic input spec (None = dynamic batch dimension)
spec = (tf.TensorSpec((None, 10), tf.float32, name="input"),)

# Convert to ONNX
model_proto, _ = tf2onnx.convert.from_keras(
    model, input_signature=spec, output_path="simple_tf_model.onnx"
)技术支持

如果您从未使用过 TensorFlow，建议先学习我们的课程：Introduction to TensorFlow in Python。

遗憾的是，最新版本的 Numpy 存在一些冲突，所以我加入了临时修复。阅读本文时，您或许已经可以不需要前六行代码即可运行。

最后，来看一下针对 scikit-learn 模型的 ONNX 转换：

from sklearn.ensemble import RandomForestClassifier
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# Train a simple model
model = RandomForestClassifier(n_estimators=10)
model.fit(X_train, y_train)

# Define input type and shape
initial_type = [('float_input', FloatTensorType([None, 4]))]

# Convert to ONNX
onnx_model = convert_sklearn(model, initial_types=initial_type)

# Save the model
with open("sklearn_model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())技术支持

与 TensorFlow 类似，DataCamp 也有免费的 Supervised Learning with scikit-learn 课程——如果上面的代码让您困惑，可以去看看。

就这些！

您现在可以运行以下代码来验证转换后的模型：

import onnx
import onnxruntime as ort
import numpy as np

# Load and check the ONNX model
onnx_model = onnx.load("simple_model.onnx")
onnx.checker.check_model(onnx_model)

# Run inference with ONNX Runtime
session = ort.InferenceSession("simple_model.onnx")
input_name = session.get_inputs()[0].name

# Create test input
test_input = np.random.randn(1, 10).astype(np.float32)

# Get prediction
onnx_output = session.run(None, {input_name: test_input})

# Compare with original model
with torch.no_grad():
    original_output = model(torch.from_numpy(test_input))

# Check if outputs match (within floating point tolerance)
np.testing.assert_allclose(
    original_output.numpy(),
    onnx_output[0],
    rtol=1e-3,
    atol=1e-5
)
print("Model conversion successful - outputs match!")技术支持

每次转换模型后都要运行此验证。小的数值差异是正常的（浮点运算并非精确），但若差异很大，说明转换过程中出了问题。

最后补充两点重要注意事项：

自定义操作无法转换：如果您编写了自定义的 PyTorch 操作或 TensorFlow 层，转换器不知道如何翻译。您需要用标准操作重写，或注册自定义的 ONNX 算子。
版本不匹配会导致问题：您的框架版本、转换器版本和 ONNX 版本必须兼容。请查看转换器文档以了解受支持的版本组合。

接下来，来聊聊 ONNX Runtime。

使用 ONNX Runtime 进行推理

本节将演示如何使用 ONNX Runtime 加载模型并进行预测。我会介绍硬件加速选项，并说明 ONNX Runtime 在生产中的优势。

ONNX Runtime 简介

ONNX Runtime 是跨平台推理引擎。微软开发它以在任何硬件上高速运行 ONNX 模型——CPU、GPU、移动芯片和专用 AI 加速器。

其工作方式如下：您加载一个 ONNX 模型，ONNX Runtime 构建执行计划。它会分析计算图、应用优化，并找出在您的硬件上执行操作的最佳方式。随后根据优化后的计划运行推理。

该架构将执行逻辑与硬件特定代码解耦。执行提供程序负责硬件接口。想在 NVIDIA GPU 上运行？使用 CUDA 执行提供程序。想用 Apple Silicon？底层代码无需改变。

ONNX Runtime 具有以下执行提供程序：

CPUExecutionProvider：默认提供程序，适用所有环境
CUDAExecutionProvider：NVIDIA GPU 上的 CUDA
TensorRTExecutionProvider：NVIDIA GPU 上的 TensorRT 优化
CoreMLExecutionProvider：Apple 设备（macOS、iOS）
DmlExecutionProvider：Windows 的 DirectML（适配任意 GPU）
OpenVINOExecutionProvider：英特尔 CPU 和 GPU
ROCMExecutionProvider：AMD GPU

每个提供程序会将 ONNX 操作翻译为硬件特定的指令。CPU 提供程序使用标准 CPU 操作。CUDA 提供程序使用 cuDNN 和 cuBLAS。TensorRT 提供程序会将您的模型编译为优化的 GPU 内核，等等。

模型的设置与执行

三行代码即可加载 ONNX 模型：

import onnxruntime as ort

# Create an inference session
session = ort.InferenceSession("simple_model.onnx")

# Get input and output names
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name技术支持

InferenceSession 会完成所有繁重工作。它加载模型、校验计算图并为推理做好准备。默认使用 CPU 执行提供程序。

现在开始推理：

import numpy as np
import onnxruntime as ort

# Create an inference session
session = ort.InferenceSession("simple_model.onnx")

# Get input and output names
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Prepare input data
input_data = np.random.randn(1, 10).astype(np.float32)

# Run inference
outputs = session.run(
    [output_name],
    {input_name: input_data},
)

# Get predictions
predictions = outputs[0]
print(predictions.shape)技术支持

run() 方法接收两个参数。第一个是输出名称列表（或 None 表示所有输出）。第二个是从输入名称到 numpy 数组的映射字典。

输入形状必须匹配。如果模型期望 [1, 10] 而您传入 [1, 10, 3]，推理将失败。

如需 GPU 加速，只需在创建会话时指定执行提供程序：

# For NVIDIA GPUs with CUDA
session = ort.InferenceSession(
    "model.onnx",
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)

# For Apple Silicon
session = ort.InferenceSession(
    "model.onnx",
    providers=['CoreMLExecutionProvider', 'CPUExecutionProvider']
)

# For Intel hardware
session = ort.InferenceSession(
    "model.onnx",
    providers=['OpenVINOExecutionProvider', 'CPUExecutionProvider']
)技术支持

按优先级列出提供程序。ONNX Runtime 会先尝试第一个，如果不可用则回退到下一个。始终将 CPUExecutionProvider 作为最后的回退选项。

您可以检查实际使用的是哪个提供程序：

import numpy as np
import onnxruntime as ort

# Create an inference session
session = ort.InferenceSession(
    "simple_model.onnx", providers=["CoreMLExecutionProvider", "CPUExecutionProvider"]
)

# Get input and output names
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Prepare input data
input_data = np.random.randn(1, 10).astype(np.float32)

# Run inference
outputs = session.run(
    [output_name],
    {input_name: input_data},
)

# Get predictions
predictions = outputs[0]
print(predictions.shape)
print("Providers:")
print(session.get_providers())技术支持

如果在期望使用 GPU 时看到 ["CPUExecutionProvider"]，说明 GPU 提供程序不可用。常见原因包括缺少驱动、ONNX Runtime 包选择错误或 GPU 不兼容。

配置不止于此。您还可以进一步调优执行提供程序以获得更好性能：

providers = [
    (
        "TensorRTExecutionProvider",
        {
            "device_id": 0,
            "trt_max_workspace_size": 2147483648,  # 2GB
            "trt_fp16_enable": True,  # Use FP16 precision
        },
    ),
    "CUDAExecutionProvider",
    "CPUExecutionProvider",
]

session = ort.InferenceSession("simple_model.onnx", providers=providers)技术支持

您可以参考官方 ONNX Runtime 文档，获取所选提供程序的全部可用选项列表。

ONNX Runtime 的应用

ONNX Runtime 的强大之处在于它几乎无所不在——服务器、浏览器、移动设备均可运行。

ONNX Runtime Web 通过 WebAssembly 和 WebGL 将推理带到浏览器。您的模型在客户端运行，无需将数据发往服务器。这对医疗影像分析、文档处理等注重隐私的应用非常适用。

以下是一个在简单 HTML 文件中使用 ONNX Runtime 的示例：

<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  <title>ONNXRuntime Web Demo</title>
</head>
<body>
  <h1>ONNXRuntime Web Demo</h1>

  <!-- Load ONNXRuntime-Web from CDN -->
  <script src="<https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js>"></script>

  <script>
    async function runModel() {
      // Create inference session
      const session = await ort.InferenceSession.create('./simple_model.onnx');

      // Prepare input
      const inputData = new Float32Array(10).fill(0).map(() => Math.random());
      const tensor = new ort.Tensor('float32', inputData, [1, 10]);

      // Run inference
      const results = await session.run({ input: tensor });

      const output = results.output;
      console.log('Output shape:', output.dims);
      console.log('Output values:', output.data);
    }

    runModel();
  </script>
</body>
</html>技术支持

对于移动端部署，您可以使用 ONNX Runtime Mobile——这是为 iOS 与 Android 优化的轻量版本。它去除了不必要的特性，减小了二进制体积。

以下为 iOS 配置示例：

import onnxruntime_objc

let modelPath = Bundle.main.path(forResource: "simple_model", ofType: "onnx")!
let session = try ORTSession(modelPath: modelPath)

let inputData = // Your input as Data
let inputTensor = try ORTValue(tensorData: inputData,
                               elementType: .float,
                               shape: [1, 10])

let outputs = try session.run(withInputs: ["input": inputTensor])技术支持

或者，如果您使用 Python 构建 REST API，可以按如下方式在某个端点中集成 ONNX Runtime：

from fastapi import FastAPI
import onnxruntime as ort
import numpy as np

app = FastAPI()
session = ort.InferenceSession("simple_model.onnx")

@app.post("/predict")
async def predict(data: dict):
    input_data = np.array(data["input"]).astype(np.float32)
    outputs = session.run(None, {"input": input_data})
    return {"prediction": outputs[0].tolist()}技术支持

就这样！其他环境的做法也大同小异，关键点是——ONNX Runtime 易于使用。接下来聊聊优化。

优化 ONNX 模型

本节将演示如何通过量化与图优化让 ONNX 模型更快、更小。我会介绍在真实部署中最重要的技术。

量化技术

量化通过降低模型精度来实现更快、更小。

将权重由 32 位浮点存储改为 8 位整数存储。这样内存占用可减少 75%，而且在大多数硬件上整数运算比浮点运算更快，从而提升推理速度。

动态量化会将权重量化为低精度，但保持激活值（推理过程中的中间值）为全精度。它是最简单的量化方法，因为不需要校准数据。

以下是应用动态量化的方法：

from onnxruntime.quantization import quantize_dynamic, QuantType

# Quantize the model
quantize_dynamic(
    model_input="model.onnx",
    model_output="model_quantized.onnx",
    weight_type=QuantType.QUInt8  # 8-bit unsigned integers
)技术支持

就是这样。量化后的模型现在使用 8 位无符号整数，已经可以直接使用。

静态量化会将权重和激活值都转换为低精度。它比动态量化更激进、速度更快，但需要具有代表性的校准数据来衡量激活范围。

下面是应用静态量化的一般策略：

from onnxruntime.quantization import quantize_static, CalibrationDataReader
import numpy as np

# Create a calibration data reader
class DataReader(CalibrationDataReader):
    def __init__(self, calibration_data):
        self.data = calibration_data
        self.index = 0

    def get_next(self):
        if self.index >= len(self.data):
            return None
        batch = {"input": self.data[self.index]}
        self.index += 1
        return batch

# Load calibration data (100-1000 samples from your dataset)
calibration_data = [np.random.randn(1, 10).astype(np.float32)
                    for _ in range(100)]
data_reader = DataReader(calibration_data)

# Quantize
quantize_static(
    model_input="model.onnx",
    model_output="model_static_quantized.onnx",
    calibration_data_reader=data_reader
)技术支持

校准数据应能代表您的实际推理负载，因此不要像上面演示那样使用随机数据。请使用来自您数据集的真实样本。

总体而言，作为机器学习工程师，您需要了解量化的收益与权衡。以下是收益：

模型体积缩小 4 倍：INT8 的内存占用是 FP32 的四分之一
推理速度提升 2-4 倍：CPU 上整数运算更快
更低的能耗：计算更少，能耗更低
更好的缓存利用：更小的模型更容易装入 CPU 缓存

以下是权衡：

轻微精度损失：INT8 通常损失 1-2%，多数场景可接受
并非所有操作都适合量化：批归一化和 softmax 比较棘手
依赖硬件：GPU 加速效果不一，部分 GPU 不支持 INT8 加速

量化对于大语言模型和生成式 AI 也很适用。一个 70 亿参数模型在 FP32 下需要 28GB 内存；量化为 INT8 后降至 7GB；量化为 INT4 后是 3.5GB——足够在消费级硬件上运行。

图优化

图优化通过重写模型计算图来提速。

节点融合会将多个操作合并为单一操作。如果模型中有 Conv 层后接 BatchNorm 再接 ReLU，优化器会将它们融合为一个 ConvBnRelu 节点。这样可以减少内存传输并加快执行。

融合前的计算图如下：

Input -> Conv -> BatchNorm -> ReLU -> Output技术支持

融合后被简化为：

Input -> ConvBnRelu -> Output技术支持

谈到图优化，您还需要了解以下两个概念：

常量折叠：对不依赖输入数据的操作进行预计算。如果模型将权重矩阵与常量相乘，优化器会在优化阶段执行一次该乘法，而不是每次推理都计算。
死代码消除：移除未使用的操作。如果模型中存在永不执行的分支或永不使用的输出，优化器会将其剔除。

以下是在 ONNX Runtime 中启用图优化的方法：

import onnxruntime as ort

# Set optimization level
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

# Create session with optimizations
session = ort.InferenceSession(
    "simple_model.onnx",
    sess_options,
    providers=["CPUExecutionProvider"]
)技术支持

ONNX Runtime 提供以下优化级别：

ORT_DISABLE_ALL：不进行优化（适合调试）
ORT_ENABLE_BASIC：不会改变数值结果的安全优化
ORT_ENABLE_EXTENDED：更激进的优化，可能带来轻微数值差异
ORT_ENABLE_ALL：所有优化，包括布局变换

生产环境建议使用 ORT_ENABLE_ALL。相较微小的数值差异，性能提升更值得。

您还可以运行以下代码，将优化后的模型保存以便复用：

sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
sess_options.optimized_model_filepath = "model_optimized.onnx"

# This creates and saves the optimized model
session = ort.InferenceSession("simple_model.onnx", sess_options)技术支持

图优化不受模型来源框架的影响。PyTorch 的 Conv-BN-ReLU 模式与 TensorFlow 的 Conv-BN-ReLU 模式会以相同方式被融合。这正是共享优化的好处——优化逻辑编写一次，即可用于任意框架生成的模型。

接下来，进入部署话题。

部署场景

本节将介绍三种部署场景——边缘设备、云端基础设施和网页浏览器。我会说明每种环境的挑战与对应的解决方案。

边缘设备部署

边缘设备的核心挑战在于资源有限。

智能手机通常只有 4–16GB 内存。树莓派更少。物联网设备可能只有 512MB 或更低。您不能指望把 70 亿参数的模型直接丢给这些设备就能运行。

从量化开始。 INT8 模型应作为基线。如果能接受精度损失，可进一步用 INT4。这样可使模型足够小以适配内存。

在边缘设备上部署时，ONNX Runtime Mobile 是好帮手。它去除了服务器端用不到的特性，并针对电池续航进行了优化。二进制更小、启动更快、功耗更低。

如下将其添加到移动应用项目：

# iOS
pod 'onnxruntime-mobile-objc'

# Android
implementation 'com.microsoft.onnxruntime:onnxruntime-android:latest.version'技术支持

执行提供程序处理硬件差异。Android 设备使用 ARM CPU，有的配备来自不同厂商的 GPU，更新的手机还有 NPU（神经网络处理单元）。您不想为每颗芯片分别写代码。

参考以下代码选择合适的执行提供程序：

# iOS - use CoreML for Apple Silicon optimization
session = ort.InferenceSession(
    "simple_model.onnx",
    providers=["CoreMLExecutionProvider", "CPUExecutionProvider"]
)

# Android - use NNAPI for hardware acceleration
session = ort.InferenceSession(
    "simple_model.onnx",
    providers=["NnapiExecutionProvider", "CPUExecutionProvider"]
)技术支持

iOS 上的 CoreML 会使用神经引擎（Apple 的 NPU）。Android 上的 NNAPI 会利用设备所具备的任意加速器——GPU、DSP 或 NPU。提供程序抽象了硬件，因此您的代码保持不变。新芯片发布且性能更好时，只需更新执行提供程序，您的应用无需改代码即可更快。

云端部署

云端部署带来无限扩展与最新硬件。

您可以在任意云平台部署 ONNX 模型——AWS、Azure、GCP，或自建服务器。模式相同——将推理服务容器化，并部署到负载均衡器后面。

以下是一个小型 Python FastAPI 示例：

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import onnxruntime as ort
import numpy as np
import logging

app = FastAPI()
logger = logging.getLogger(__name__)

# Load model at startup
session = None

@app.on_event("startup")
async def load_model():
    global session
    session = ort.InferenceSession(
        "simple_model.onnx",
        providers=[
            "CUDAExecutionProvider",
            "CoreMLExecutionProvider",
            "CPUExecutionProvider",
        ],
    )
    logger.info("Model loaded successfully")

class PredictionRequest(BaseModel):
    input: list

class PredictionResponse(BaseModel):
    prediction: list
    inference_time_ms: float

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    try:
        import time

        start = time.time()

        # Input shape: (1, 10)
        input_data = np.array(request.input, dtype=np.float32).reshape(1, 10)
        input_name = session.get_inputs()[0].name
        outputs = session.run(None, {input_name: input_data})

        inference_time = (time.time() - start) * 1000

        return PredictionResponse(
            prediction=outputs[0].tolist(), inference_time_ms=inference_time
        )
    except Exception as e:
        logger.error(f"Prediction failed: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health():
    return {"status": "healthy", "model_loaded": session is not None}技术支持

这个 FastAPI 示例虽小，但对机器学习工程师而言不一定最直观。您可通过我们的免费 Introduction to FastAPI 课程掌握基础。

随后用 Docker 将其容器化：

FROM python:3.13-slim

WORKDIR /app

# Install dependencies
RUN pip install onnx onnxruntime onnxscript fastapi pydantic numpy uvicorn

# Copy model and code
COPY simple_model.onnx .
COPY main.py .

# Expose port
EXPOSE 8000

# Run the service
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]技术支持

每位机器学习爱好者都应掌握 Docker 基础。参加我们的 Introduction to Docker 课程，学习如何部署机器学习模型。

应用将运行在 8000 端口，您可以直接打开文档并试用预测端点：

如果您使用 Azure Machine Learning，会很高兴得知它可直接集成 ONNX。三步即可部署模型：

from azureml.core import Workspace, Model
from azureml.core.webservice import AciWebservice, Webservice

# Register model
ws = Workspace.from_config()
model = Model.register(
    workspace=ws, model_path="simple_model.onnx", model_name="my-onnx-model"
)

# Deploy
deployment_config = AciWebservice.deploy_configuration(cpu_cores=1, memory_gb=1)

service = Model.deploy(
    workspace=ws,
    name="onnx-service",
    models=[model],
    deployment_config=deployment_config,
)技术支持

Azure 会处理扩缩容、监控与更新。您可获得自动日志、请求跟踪与健康检查。

若您是 Azure Machine Learning 新手，请阅读我们的初学者指南，快速了解基础。

谈到 MLOps 工作流，需要以下组件：

版本控制：使用 DVC 或 MLflow 跟踪模型版本
CI/CD：使用 GitHub Actions 或 Azure Pipelines 自动化测试与部署
监控：跟踪推理延迟、吞吐量与错误率
A/B 测试：在生产中比较模型版本
回滚：出现问题时回退到先前版本

以下是可用的示例代码：

import prometheus_client as prom
import onnxruntime as ort
import numpy as np

# Define metrics
inference_duration = prom.Histogram(
    "model_inference_duration_seconds", "Time spent processing inference"
)

inference_count = prom.Counter("model_inference_total", "Total number of inferences")

# Load model
session = ort.InferenceSession("simple_model.onnx")

@app.post("/predict")
@inference_duration.time()
async def predict(request: PredictionRequest):
    inference_count.inc()
    # Input shape: (1, 10)
    input_data = np.array(request.input, dtype=np.float32).reshape(1, 10)
    input_name = session.get_inputs()[0].name
    outputs = session.run(None, {input_name: input_data})
    return {"prediction": outputs[0].tolist()}技术支持

在此基础上，您可以通过在负载均衡器后运行多个服务实例实现横向扩展。ONNX Runtime 是无状态的，因此任意实例都可处理任意请求。

基于浏览器的推理

ONNX Runtime Web 借助 WebAssembly 和 WebGL 直接在浏览器中运行模型。

与其将数据发送到服务器，不如一次性把模型发送到浏览器并在本地运行推理。用户能即时获得响应，无需网络延迟。

开始很简单，新建一个包含内联 JS 脚本的 HTML 文件：

<!doctype html>
<html>
  <head>
    <meta charset="utf-8" />
    <script src="<https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js>"></script>
  </head>
  <body>
    <script>
      async function run() {
        try {
          const session = await ort.InferenceSession.create("simple_model.onnx", {
            executionProviders: ["wasm"]
          });
          console.log("Session loaded:", session);

          // Inspect names
          console.log("Inputs:", session.inputNames);
          console.log("Outputs:", session.outputNames);

          const inputName = session.inputNames[0];

          const inputData = new Float32Array(10).fill(0).map(() => Math.random());
          const tensor = new ort.Tensor("float32", inputData, [1, 10]);

          const feeds = {};
          feeds[inputName] = tensor;

          const results = await session.run(feeds);

          const outputName = session.outputNames[0];
          console.log("Prediction:", results[outputName].data);

        } catch (err) {
          console.error("ERROR:", err);
        }
      }

      run();
    </script>
  </body>
</html>技术支持

使用 live server 打开文件即可在控制台看到输出：

最后，我们来讨论几个高级话题。

高级话题

标准 ONNX 操作覆盖了大多数用例，但如果您需要更多功能，请看这一节。

我将介绍自定义算子和大语言模型部署。会解释在何种情况下需要这些高级特性，以及它们带来的权衡。

自定义算子

ONNX 自带数百个操作，但您的模型可能使用了独特的东西。

自定义算子允许您定义标准 ONNX 算子集中不存在的操作。您自行编写操作逻辑，并告知 ONNX Runtime 如何执行。这样即可扩展 ONNX 的内置能力。

何时需要自定义算子？当您使用 ONNX 尚未支持的前沿研究成果时；或当您构建了特定领域的专有操作时；或当您需要标准算子无法提供的硬件特定优化时。

但问题在于——自定义算子会破坏可移植性。没有您的自定义实现，模型无法运行。如果您将包含自定义算子的 PyTorch 模型导出到 ONNX，则需要单独打包这些算子，并在 ONNX Runtime 中注册。

部署也会更难。每个环境都需要您的自定义算子库。边缘设备、云服务器、浏览器都需要相同的实现。您也失去了 ONNX 最大的优势：一次编写，到处运行。

以下是一个自定义算子的简单示例，便于您了解其工作方式：

import numpy as np
import onnxruntime as ort
from onnxruntime import InferenceSession, SessionOptions

# Define the custom operation
def custom_square(x):
    return x * x

# Register the custom op with ONNX Runtime
# This tells ONNX Runtime how to execute "CustomSquare" operations
class CustomSquareOp:
    @staticmethod
    def forward(x):
        return np.square(x).astype(np.float32)

# Create a simple ONNX model with a custom op using the helper
from onnx import helper, TensorProto
import onnx

# Create a graph with custom operator
node = helper.make_node(
    "CustomSquare",  # Custom operation name
    ["input"],
    ["output"],
    domain="custom.ops",  # Custom domain
)

graph = helper.make_graph(
    [node],
    "custom_op_model",
    [helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 10])],
    [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 10])],
)

model = helper.make_model(graph)
onnx.save(model, "custom_op_model.onnx")

# To use this model, you'd need to implement the custom op in C++
# and register it with ONNX Runtime - Python-only custom ops aren't supported
print("Custom operator model created - requires C++ implementation to run")技术支持

多数团队会尽量避免使用自定义算子。优先尝试用标准算子的组合来实现。只有当标准算子无法表达您的操作，或性能上确有必要时，才考虑自定义。

处理大语言模型

LLM 的主要问题是内存。一个 70 亿参数模型在 FP32 需要 28GB、在 FP16 需要 14GB、在 INT8 需要 7GB。大多数消费级硬件承载不了。您需要量化，且往往需要激进的 INT4 或 INT3 量化。

如果您是 LLM 新手并想进一步了解，我们的 Large Language Models (LLMs) Concepts 课程是很好的起点。

KV 缓存管理对自回归生成很重要。生成每个 token 都需要引用此前所有 token。缓存会随序列长度增长。长对话会快速吃掉内存。ONNX Runtime GenAI 包含优化的 KV 缓存管理，可复用内存并降低开销。

推理速度取决于硬件与模型大小。较小的模型（小于 30 亿参数）在消费级 GPU 上表现良好。更大的模型需要多块 GPU 或专用推理加速器。ONNX Runtime 通过执行提供程序支持多 GPU 推理，但您需要手动将模型拆分到多设备上。

Flash Attention 可加速 Transformer 中的注意力机制。标准注意力在序列长度上是 O(n²)——增长很快。Flash Attention 通过减少内存搬运提升速度且不改变结果。ONNX Runtime 在受支持的硬件上包含 Flash Attention 优化。

用 ONNX Runtime 训练 LLM 是可能的，但并不常见。多数团队使用 PyTorch 或 JAX 训练，然后导出为 ONNX 仅用于推理。ONNX Runtime Training 适用于微调与继续预训练，但围绕 PyTorch 训练的生态更成熟。

生成式 AI 工作负载在 ONNX Runtime 中享有特殊支持。GenAI 扩展提供用于文本生成的高级 API，能在调用间维护状态，并管理束搜索或采样策略。这样您就无需自行实现生成逻辑。

LLM 领域变化迅速，ONNX Runtime 也在尽力跟进。算子集与优化会定期更新，但研究到生产级支持之间总会有一定滞后。

结语

现在，您已经掌握了如何从任意框架转换模型、为生产进行优化，并在任何地方部署。

ONNX 消除了训练与部署之间的摩擦。用擅长科研的 PyTorch 进行训练；用快速且“无处不在”的 ONNX Runtime 部署。量化到 INT8，将推理成本减半。无需改代码即可从 CPU 切换到 GPU。在不同云服务商之间迁移模型而不被锁定。

ONNX 让这一切成为可能。

生态系统也在持续进步。生成式 AI 支持随每个版本更完善——更快的 LLM 推理、更好的内存管理，以及新的优化技术。硬件支持扩展至更多加速器与边缘设备。社区也在构建让 ONNX 更易用的工具。您可以通过测试新特性、报告问题，或在 GitHub 仓库贡献代码来参与其中。

准备好更进一步了吗？掌握可解释人工智能（XAI）概念，别再将模型视作黑盒。

作者

Dario Radečić

常驻克罗地亚的高级数据科学家。顶级科技写作者，已发表 700 多篇文章，累计获得超过 1,000 万次阅读。著有《使用 TPOT 实现机器学习自动化》。

常见问题

什么是 ONNX，为什么要使用它？

ONNX（Open Neural Network Exchange）是面向机器学习模型的通用格式，适用于不同框架。它让您可以在 PyTorch 或 TensorFlow 中训练模型，并在任意环境部署而无需重写代码。您将获得框架无关性、通过优化带来的更高性能，以及可在任何硬件（从边缘设备到云服务器）上运行模型的能力。

ONNX 如何提升模型性能？

ONNX Runtime 会应用图优化（如节点融合、常量折叠）来加速推理。量化可将模型大小减少最高 75%，并通过将权重从 32 位浮点转换为 8 位整数，使推理速度提升 2–4 倍。这些优化可自动发挥作用，您无需手动微调模型。

ONNX 能在移动设备和浏览器中运行吗？

可以。ONNX Runtime Mobile 可在 iOS 和 Android 设备上运行，并针对电池和内存使用做了优化。ONNX Runtime Web 则通过 WebAssembly 和 WebGL 直接在浏览器中运行模型，让您可以在客户端进行推理而无需将数据发往服务器。这两种方式都支持通过执行提供程序进行硬件加速，如 Apple 设备上的 CoreML、Android 上的 NNAPI。

动态量化与静态量化有什么区别？

动态量化将权重量化为低精度，但保持激活为全精度，并且不需要校准数据。静态量化会同时对权重和激活进行低精度转换，速度更快，但需要具有代表性的校准数据来衡量激活范围。静态量化性能更好，但设置工作更多。

将含自定义操作的模型转换为 ONNX 时应如何处理？

自定义操作需要使用 C++ 实现并在 ONNX Runtime 中注册——仅用 Python 的自定义算子不适合生产使用。您需要用 C++ 实现自定义算子，将其编译为共享库，并在每个部署环境中注册。多数团队会尽量避免自定义算子，而是使用标准 ONNX 操作的组合，因为自定义算子会破坏跨平台的可移植性。

主题

机器学习

主题

技术

ONNX: Train in Any Framework, Deploy on Any Hardware

使用 AI 探索

ONNX 入门

先决条件

环境配置

理解 ONNX 基础

ONNX 模型格式

关键概念与架构

将模型转换为 ONNX 格式

框架支持与工具

分步转换流程

使用 ONNX Runtime 进行推理

ONNX Runtime 简介

模型的设置与执行

ONNX Runtime 的应用

优化 ONNX 模型

量化技术

图优化

部署场景

边缘设备部署

云端部署

基于浏览器的推理

高级话题

自定义算子

处理大语言模型

结语

常见问题

主题

与 DataCamp 一起学习

Python 中的机器学习特征工程

机器学习科学家在 Python 中

使用 Keras 构建语言建模的循环神经网络（RNN）

主题

技术

与 DataCamp 一起学习

Python 中的机器学习特征工程

机器学习科学家在 Python 中

使用 Keras 构建语言建模的循环神经网络（RNN）

ONNX 入门

先决条件

环境配置

理解 ONNX 基础

ONNX 模型格式

关键概念与架构

将模型转换为 ONNX 格式

框架支持与工具

分步转换流程

使用 ONNX Runtime 进行推理

ONNX Runtime 简介

模型的设置与执行

ONNX Runtime 的应用

优化 ONNX 模型

量化技术

图优化

部署场景

边缘设备部署

云端部署

基于浏览器的推理

高级话题

自定义算子

处理大语言模型

结语

常见问题

主题

与 DataCamp 一起学习

Python 中的机器学习特征工程

机器学习科学家 在 Python 中

使用 Keras 构建语言建模的循环神经网络（RNN）

与 DataCamp 一起学习

Python 中的机器学习特征工程

机器学习科学家 在 Python 中

使用 Keras 构建语言建模的循环神经网络（RNN）

机器学习科学家在 Python 中

机器学习科学家在 Python 中