领域专用小型语言模型——探索 ONNX

本章涵盖：

• ONNX 标准格式
• ONNX Runtime
• ONNX 如何在有无硬件加速的情况下用于 LLM

本章介绍 ONNX 框架。它在模型优化、量化，以及跨框架和跨硬件厂商的可移植性方面发挥着重要作用。如果你刚接触 ONNX，请花些时间吸收本章概念——后续章节会大量使用它们。

5.1 ONNX 格式

ONNX，即 Open Neural Network Exchange，是一个用于机器学习互操作性的开放标准。它最早发布于 2017 年，现在是 Linux Foundation for Artificial Intelligence，也就是 LFAI 的毕业项目。ONNX 旨在提升机器学习，也就是 ML，以及深度学习，也就是 DL，框架之间的互操作性，包括 Keras、TensorFlow、PyTorch、scikit-learn、XGBoost 等；同时最大化不同硬件加速器上的性能，不只是 NVIDIA，也包括 Intel OpenVINO、Habana、Qualcomm、Apache TVM、Hugging Face Optimum 等。图 5.1 从高层展示了 ONNX 生态系统，包括与框架和平台无关的模型格式，以及可选的 ONNX Runtime。

图 5.1 —— ONNX 概览

随着 ML 和 DL 框架不断演进，可移植性变得至关重要——我们今天使用的东西，未必就是明天还会使用的东西。ONNX 是一个稳健的开放标准，它减少了对特定框架和硬件加速器的锁定，有助于确保组织的模型能够长期保持可用。

ONNX 是一个开放治理的社区项目，因此任何人都可以贡献。它最初由 Facebook 创建，很快获得了 Microsoft、IBM、Huawei、Intel、AMD、Arm 和 Qualcomm 等公司的支持。此后，它已经成长为被广泛采用的 ML/DL 推理标准。图 5.2 列出了当前支持 ONNX 的公司。

ONNX 由两个项目组成：ONNX 格式和 ONNX Runtime。本节和 5.2 节介绍 ONNX 格式，5.3 节和 5.4 节介绍 ONNX Runtime。

图 5.2 —— 支持 ONNX 倡议的组织列表

除了跨框架之外，ONNX 还为多种编程语言提供 API。在本书中，我们会使用 Python，并稍微使用一点 C/C++。图 5.3 总结了 ONNX 中所有可能的选项。

图 5.3 —— ONNX 中对 ML/DL 模型可以执行的操作

一旦你在某个 DL 框架中训练了模型，或者从别人那里接收了一个模型，你可以将它转换为 ONNX，并对其进行优化，然后从以下选项中选择：

• 直接在 ONNX Runtime 中执行它，或者在该框架支持的任何其他 runtime 中执行它
• 将它转换为另一个 DL 框架
• 通过量化进一步优化它

在后两个选项所描述的流程结束时，你可以使用任何受支持的 runtime 运行最终模型。在图 5.3 中，模型在另一个框架中运行的那个模块用不同颜色高亮，因为它超出了本书范围。从这里开始，我们将讨论如何将模型转换为 ONNX 格式，并对其进行优化和量化，同时也会引用支持该标准的其他 Python 库。

ONNX 定义了一组通用算子和一种标准文件格式，用于存储训练好的 ML/DL 模型和流水线，并提供足够细节，使它们可以在平台之间迁移。ONNX 还可以将已存储的操作编译为更底层的语言，以便嵌入到一系列设备中，包括边缘设备。从概念上说，一个 ONNX 文件定义的是一个有向图：每条边表示一个带类型的 tensor，从一个节点移动到下一个节点。节点被称为 operator，也就是算子：它们消费输入，并把结果传递给子节点。ONNX 的标准操作足够广泛，能够表达大多数 ML/AI 工作负载；必要时也可以通过自定义算子进行扩展。

ONNX 提供了一种通用语言，使任何 ML/DL 框架都能够描述自己的模型，从而让生产部署更容易。它也允许你构建一条独立于模型创建框架的统一部署流水线。该框架包含一个用于评估 ONNX 模型和算子的 Python runtime，但使用时要小心：它的用途是澄清 ONNX 语义并帮助调试，而不是用于生产环境——尤其是在你需要高性能时。

无论 ML/DL 模型或框架是什么，构建 ONNX 图都涉及定义一组算子的组合，下一节将对此进行说明。

5.2 ONNX 算子和类型

算子是定义 ONNX 模型的基本构建块；来自不同框架的几乎任何 ML/DL 模型，都可以用这些算子来描述。ONNX 算子包括标准矩阵操作、归约、图像变换、深度神经网络层和激活函数。表 5.1 提供了一份与 LLM 相关的非完整算子列表，完整列表可参考官方 ONNX 文档。

表 5.1 —— 与 LLM 最相关的 ONNX 算子

算子	描述
Add	执行逐元素二元加法。支持多方向广播，类似 NumPy。
And	对输入 tensor 逐元素执行 and 逻辑操作，并返回结果 tensor。
ArgMax	沿给定轴计算输入 tensor 元素中最大值的索引。
ArgMin	沿给定轴计算输入 tensor 元素中最小值的索引。
Atan	对给定输入 tensor 逐元素计算反正切，也就是 tangent 的反函数。
AveragePool	消费一个输入 tensor，并根据 kernel size、stride size 和 pad length 对 tensor 应用平均池化。平均池化会根据 kernel size 对输入 tensor 子集中的所有值求平均，并把数据下采样到输出 tensor，供后续处理。
BatchNormalization	对深度神经网络应用 batch normalization。
Conv	对输入 tensor 应用卷积。
DFT	计算输入的离散傅里叶变换。
Det	计算方阵，或一批方阵的行列式。
Div	对输入执行逐元素二元除法。
Equal	对输入 tensor 逐元素执行 equal 逻辑操作，并返回结果 tensor。
Exp	对给定输入 tensor 逐元素计算指数。
Floor	接收一个输入 tensor，并生成一个输出 tensor，其中逐元素应用 y = floor(x)。
GreaterOrEqual	对输入 tensor 逐元素执行 greater_equal 逻辑操作，并返回结果 tensor。
Identity	恒等算子。
IsNaN	返回输入中哪些元素是 NaN。
LayerNormalization	以函数形式定义的 layer normalization。
LessOrEqual	对输入 tensor 逐元素执行 less equal 逻辑操作，并返回结果 tensor。
Loop	通用循环结构。
MatMul	矩阵乘积。
MaxPool	消费一个输入 tensor，并根据 kernel size、stride size 和 pad length 对 tensor 应用最大池化。最大池化操作会根据 kernel size 对输入 tensor 子集中的所有值取最大值，并把数据下采样到输出 tensor，供后续处理。
Mean	对每个输入 tensor 逐元素求平均。
Pad	给定一个包含待 padding 数据的 tensor、一个包含各轴起始和结束 padding 值数量的 tensor，以及可选的 mode 和 constant value，生成一个经过 padding 的 tensor 作为输出。
QuantizeLinear	线性量化算子。它消费一个高精度 tensor、一个 scale 和一个 zero point，以计算低精度 / 量化 tensor。
Reciprocal	接收一个输入 tensor，并生成一个输出 tensor，其中 reciprocal，即 y = 1/x，逐元素应用到 tensor。
RegexFullMatch	对输入 tensor 的每个元素执行完整 regex 匹配。如果某个元素完全匹配作为 attribute 指定的 regex pattern，则输出中的对应元素为 true，否则为 false。
Reshape	重塑输入 tensor。第一个输入是 data tensor；第二个输入是 shape tensor，用于指定输出形状。它输出重塑后的 tensor。
Softmax	对给定输入计算归一化指数值：Softmax(input, axis) = Exp(input) / ReduceSum(Exp(input), axis=axis, keepdims=1)。axis attribute 表示沿哪个维度执行 Softmax。输出 tensor 与输入形状相同，并包含对应输入的 Softmax 值。
Transpose	转置输入 tensor。
Unique	查找 tensor 中的唯一元素。当提供可选 attribute axis 时，返回沿该轴切片得到的唯一 subtensor。否则，输入 tensor 会被展平，并返回展平 tensor 中的唯一值。

ONNX 针对使用 tensor 进行数值计算做了优化。tensor 是多维数组，由元素类型、形状和连续值数组定义。虽然 ONNX 最初只支持 32 位浮点值，但后续版本增加了所有常见类型。你可以用程序列出所有支持的 ONNX 类型：

import re
from onnx import TensorProto

reg = re.compile('^[0-9A-Z_]+$')

values = {}
for att in sorted(dir(TensorProto)):
    if att in {'DESCRIPTOR'}:
        continue
    if reg.match(att):
        values[getattr(TensorProto, att)] = att
for i, att in sorted(values.items()):
    si = str(i)
    if len(si) == 1:
        si = " " + si
    print("%s: onnx.TensorProto.%s" % (si, att))

前面的代码会把支持类型列表打印到标准输出：

 0: onnx.TensorProto.UNDEFINED
 1: onnx.TensorProto.FLOAT
 2: onnx.TensorProto.UINT8
 3: onnx.TensorProto.INT8
 4: onnx.TensorProto.UINT16
 5: onnx.TensorProto.INT16
 6: onnx.TensorProto.INT32
 7: onnx.TensorProto.INT64
 8: onnx.TensorProto.STRING
 9: onnx.TensorProto.BOOL
10: onnx.TensorProto.FLOAT16
11: onnx.TensorProto.DOUBLE
12: onnx.TensorProto.UINT32
13: onnx.TensorProto.UINT64
14: onnx.TensorProto.COMPLEX64
15: onnx.TensorProto.COMPLEX128
16: onnx.TensorProto.BFLOAT16
17: onnx.TensorProto.FLOAT8E4M3FN
18: onnx.TensorProto.FLOAT8E4M3FNUZ
19: onnx.TensorProto.FLOAT8E5M2
20: onnx.TensorProto.FLOAT8E5M2FNUZ

作为使用 ONNX API 的具体示例，我们来看一个线性回归模型，它由如下表达式表示：

Y = XA + B

所有变量都是矩阵。在 ONNX 中，我们必须用两个算子表示它：MatMul 和 Add，见表 5.1 中的算子列表：

y = Add(MatMul(X, A), B)

为了构建图，我们会使用 ONNX API 中的四个函数：

• make_tensor_value_info——声明一个具有给定形状和类型的输入或输出变量。
• make_node——根据一个操作及其输入和输出创建节点。
• make_graph——根据前两个函数创建的对象创建 ONNX 图。
• make_model——将图与额外 metadata 合并。

ONNX 是强类型的，因此必须定义每个函数输入和输出的形状和类型。ONNX 构建在 protobuf 之上。protobuf 是一种语言和平台中立的格式，用于序列化结构化数据。protobuf 提供描述 ML/DL 模型所需的定义；ONNX 则用于序列化和反序列化它。protobuf 负责保持已保存图的紧凑性。

要用代码实现这个例子，使用 Python binding，我们首先需要导入所需类和函数：

from onnx import TensorProto
from onnx.helper import (
    make_model, make_node, make_graph,
    make_tensor_value_info)
from onnx.checker import check_model

导入的函数包括前面列出的那些，另外还有 check_model，它用于检查所创建 ONNX 模型的一致性。TensorProto 是 ONNX 中定义 tensor 的类，也是多个基于 protobuf 的结构类之一。你可以通过如下方式用程序列出所有这些类：

import onnx
import pprint

pprint.pprint([p for p in dir(onnx)
               if p.endswith('Proto') and p[0] != '_'])

它会返回如下输出：

['AttributeProto', 'FunctionProto', 'GraphProto', 'MapProto', 'ModelProto', 'NodeProto', 'OperatorProto', 'OperatorSetIdProto', 'OperatorSetProto', 'OptionalProto', 'SequenceProto', 'SparseTensorProto', 'StringStringEntryProto', 'TensorProto', 'TensorShapeProto', 'TrainingInfoProto', 'TypeProto', 'ValueInfoProto']

我们可以先声明相关矩阵，从而定义一个线性回归模型：

X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])

接下来，创建图节点，也就是该模型实现中的操作：

node1 = make_node('MatMul', ['X', 'A'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])

然后创建图：

graph = make_graph([node1, node2], 'lr', [X, A, B], [Y])
onnx_model = make_model(graph)

按如下方式检查模型一致性：

check_model(onnx_model)

最后，可以用 SerializeToString 方法将其序列化：

with open("linear_regression.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

我们可以用两种方式检查 ONNX 模型结构。第一种，把模型变量打印到屏幕上：

print(onnx_model)

该命令会以 ASCII 格式把模型结构打印到标准输出，如下所示：

ir_version: 9
graph {
  node {
    input: "X"
    input: "A"
    output: "XA"
    op_type: "MatMul"
  }
  node {
    input: "XA"
    input: "B"
    output: "Y"
    op_type: "Add"
  }
  name: "lr"
  input {
    name: "X"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
          }
          dim {
          }
        }
      }
    }
  }
  input {
    name: "A"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
          }
          dim {
          }
        }
      }
    }
  }
  input {
    name: "B"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
          }
        }
      }
    }
  }
  output {
    name: "Y"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
          }
        }
      }
    }
  }
}
opset_import {
  version: 21
}

另一种方式是使用开源 Netron 库可视化 ONNX 图，Netron 支持 ONNX。输出会类似图 5.4，也可以在渲染中包含 metadata。

图 5.4 —— 使用 Netron 可视化 ONNX 中的线性回归模型

已保存的 ONNX 图可以随时使用 load 函数恢复：

from onnx import load

with open("linear_regression.onnx", "rb") as f:
    onnx_model = load(f)

5.3 ONNX Runtime

ONNX 格式定义了转换到该框架中的 ML/DL 模型。生态系统中的另一个关键部分是 ONNX Runtime，也就是 ORT。它可以在广泛的硬件和软件平台上高效执行 ONNX 模型，前面的图 5.1 已经展示过。ORT 支持在本地服务器、云服务器，以及手机和边缘设备上部署 ONNX 模型，从而在不同环境中保持一致性能。当硬件加速可用时，ORT 会使用硬件加速来实现更快推理。它还包含 kernel fusion、quantization 和 layout transformation 等优化，以便在保持良好准确率的同时加速执行。虽然 ORT 同时支持训练和推理，但它最常用于推理，而训练通常交给原始 ML/DL 框架完成。使用 ORT 训练超出了本书范围。

注意
ORT 提供多种编程语言的 binding；本书中我们会使用 Python。

推理过程遵循一组固定步骤。一旦我们训练了模型并把它转换为 ONNX，或者直接使用 ONNX API 实现了一个模型，例如 5.2 节中的线性回归模型，我们就可以加载它并创建一个 inference session，使用 ORT 的 InferenceSession 类：

import onnxruntime as rt

onnx_model_path = "path_to_the_onnx_model"
session = rt.InferenceSession(onnx_model_path)

接下来准备模型输入，在这个示例中，它将是一个 NumPy 数组：

import numpy as np

input_data = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)

现在可以运行推理：

input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: input_data})

最后打印结果：

print(f"Output: {result[0]}")

向模型传入输入值时，要确保使用正确的输入名称。你可以通过 inference session 对象的 get_inputs 获取它，该方法返回模型输入名称列表。类似地，get_outputs 返回模型输出名称。ORT 推理的命令序列总是相同的；随架构变化的是输入预处理、输出后处理，以及如何处理多个输入和输出。

还有一个用于在 Web 浏览器或 Node.js 上运行 ONNX 模型的 JavaScript 库，叫 ONNX Runtime Web。它构建在 WebAssembly 和 WebGL 之上。它支持 CPU 和 GPU 执行，并支持交互式、免安装、设备无关的 ML，同时降低服务端—客户端延迟，并改善隐私与安全性。当你专门化模型、将其转换为 ONNX 并完成优化之后，可以把它作为另一种 serving 选项。

5.4 ONNX Runtime Provider

为了充分利用 ONNX 模型所运行的底层硬件，ORT 使用 runtime provider，也叫 execution provider。这些接口允许你跨环境部署 ONNX 模型，并通过使用每个平台的计算能力来优化执行。通过专用 execution provider，ORT 会将特定节点或子图分配到部署环境中可用的硬件、驱动程序和相关库上运行。图 5.5 来自官方 ORT 文档，展示了 ORT 如何针对给定硬件规格改善模型执行。

图 5.5 —— ONNX Runtime 执行模型

截至本文写作时，ORT 提供了许多 execution provider。有些已经生产可用；另一些是用于评估的 preview release。标记为 in review 的 provider 通常较稳定，可能适合生产环境，但在使用任何仍标记为 preview 的 provider 部署 ONNX 模型之前，一定要仔细测试。

下面是 ORT 当前稳定的 execution provider：

• Default CPU——厂商无关 CPU。
• Intel OpenVINO——在 CPU 和 GPU 上使用 OpenVINO toolkit，包括边缘和移动设备。
• Intel oneDNN——使用 Intel Math Kernel Library for Deep Neural Networks，也就是 Intel DNNL 中的优化 primitive。仅 CPU。
• XNNPACK——使用 XNNPACK，这是一个针对 ARM、Android/iOS、WebAssembly 和 x86 优化的浮点神经网络推理库。
• NVIDIA CUDA——在支持 CUDA 的 NVIDIA GPU 上启用硬件加速计算。
• NVIDIA TensorRT——在支持的 NVIDIA GPU 上使用 NVIDIA 的 TensorRT engine。
• DirectML——使用 DirectML 的 ML 低层 API；仅 Windows GPU。
• AMD MIGraphX——AMD 面向 AMD GPU 的 MiGraphX 优化引擎。
• AMD ROCm——用于带 ROCm 的 AMD GPU。
• Android Neural Networks API——面向 Android 设备上的 CPU、GPU 和神经网络加速器的统一接口。
• Qualcomm QNN——在 Qualcomm 芯片组上使用 Qualcomm AI Engine Direct SDK。

下面是当前仍处于 preview 状态的 ORT execution provider 列表：

• TVM——使用 Apache TVM，为 CPU、GPU、边缘和移动设备优化模型。目前已在 Linux 和 Windows 上测试，但尚未在 macOS 上测试。
• ARM Compute Library——使用 Arm Compute Library engine 在 Armv8 core 上运行推理。
• ARM-NN——面向 Armv8 core 的 ArmNN inference engine。
• CoreML——使用 Core ML framework，在 Apple 的 CPU、GPU 和 Neural Engine 上运行。
• Rockchip-NPU——使用 RKNPU DDK。
• Xilinx Vitis-AI——在 AMD 平台上使用 AMD 的开发栈进行硬件加速 AI 推理。
• Huawei CANN——在 Huawei Ascend 硬件上使用 CANN Execution Provider。
• AZURE——调用远程 Azure endpoint 执行推理。

你可以通过程序检查目标环境中可用的 provider：

import onnxruntime as rt

rt.get_available_providers()

这会返回一个 Python 列表。例如，在一个带 CPU 和 NVIDIA GPU，并且已安装 CUDA 驱动的系统上，前面的方法会返回：

['CPUExecutionProvider', 'CUDAExecutionProvider']

如果有多个 provider 可用，你可以告诉 ORT 按什么顺序尝试它们。示例如下：

providers_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']

sess = rt.InferenceSession("model.onnx", providers=providers_list)

这里，我们告诉 ORT 在执行环境中可用时，优先使用 CUDA provider。

5.5 在 CPU 上将 ONNX 用于 LLM

现在我们已经介绍了 ONNX 格式和 runtime 的主要概念，是时候看看如何将它们用于 LLM 了。本节聚焦 CPU-only 设置；下一节介绍 GPU 可用时的相同思路。本节代码包含在配套 Colab notebook 中。

第一个 LLM 示例使用基于 BERT 的模型 bert-base-uncased。它使用 Hugging Face 的 Transformers 和 Datasets 库。执行环境必须安装 onnx 和 CPU-only 的 onnxruntime Python 包：

pip install onnx onnxruntime

接下来，我们需要从 Hugging Face Hub 下载模型：

from transformers import AutoModelForQuestionAnswering, BertTokenizer

model_id = 'google-bert/bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_id)
model = AutoModelForQuestionAnswering.from_pretrained(model_id)
model.eval()

为了测试模型，我们会使用 SQuAD 数据集，也就是 Stanford Question Answering Dataset 的一个子集。该数据集可以通过 Datasets 库从 Hugging Face Hub 获取：

from datasets import load_dataset

samples_count = 200
squad = load_dataset("squad", split="validation[:"+ str(samples_count) +"]")

只会使用 200 个样本。下面的示例展示了它们的结构：

{'id': '56be4db0acb8001400a502ec',
 'title': 'Super_Bowl_50',
 'context': 'Super Bowl 50 was an American football game to determine the
 champion of the National Football League (NFL) for the 2015 season. The 
American Football Conference (AFC) champion Denver Broncos defeated the 
National Football Conference (NFC) champion Carolina Panthers 24–10 to earn 
their third Super Bowl title. The game was played on February 7, 2016, at 
Levi's Stadium in the San Francisco Bay Area at Santa Clara, California. 
As this was the 50th Super Bowl, the league emphasized the "golden 
anniversary" with various gold-themed initiatives, as well as temporarily 
suspending the tradition of naming each Super Bowl game with Roman numerals 
(under which the game would have been known as "Super Bowl L"), so that the 
logo could prominently feature the Arabic numerals 50.',
 'question': 'Which NFL team represented the AFC at Super Bowl 50?',
 'answers': {'text': ['Denver Broncos', 'Denver Broncos', 'Denver Broncos'],
  'answer_start': [177, 177, 177]}}

在这个示例中，我们会使用 context 和 question 字段，对不同模型版本的答案生成进行 benchmark。

先创建一个目录，用来存放 ONNX 转换后的模型：

import os

output_dir = os.path.join(".", "onnx_models")
onnx_model_filename = 'bert-base-uncased.onnx'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
export_model_path = os.path.join(output_dir, onnx_model_filename)

在转换流程开始时，我们会使用一个测试样本运行一次模型，这里使用列表中的第一个样本，不过取回的 200 个样本中任何一个都可以：

tokenized_inputs = tokenizer(squad["question"][0],
                 ➥squad["context"][0], return_tensors="pt")
inputs = {
        'input_ids':  tokenized_inputs['input_ids'],
        'input_mask': tokenized_inputs['attention_mask'],
        'segment_ids': tokenized_inputs['token_type_ids']
    }

然后开始转换：

import torch

with torch.no_grad():
    symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}
    torch.onnx.export(model,
                      args=tuple(inputs.values()),
                      f=export_model_path,
                      opset_version=15,
                      do_constant_folding=True,
                      input_names=['input_ids',
                                       'input_mask',
                                       'segment_ids'],
                      output_names=['start', 'end'],
                      dynamic_axes={'input_ids': symbolic_names,
                                    'input_mask' : symbolic_names,
                                    'segment_ids' : symbolic_names,
                                    'start' : symbolic_names,
                                    'end' : symbolic_names})

我们会使用 PyTorch 的 ONNX export 函数，传入要转换的模型、目标目录、ONNX 版本，也就是本文写作时的 1.15.0，以及模型的输入和输出名称。我们还会启用 constant folding 进行优化，并设置可变长度轴，使模型可以在不同 batch size 和不同输入、输出 token 数量下运行。图 5.6 展示了转换后模型 ONNX 图的一部分，为了可读性进行了裁剪。

图 5.6 —— 转换后的 bert-base-uncased 模型 ONNX 图的一部分，一个 attention block

你可以使用 5.1 节介绍的 check_model 函数验证导出的模型：

from onnx.checker import check_model

check_model(export_model_path, full_check=True) 

如果没有错误信息，则说明模型 schema 有效。要验证完整架构，可以用原始模型和 ONNX 转换后模型在若干样本，或者完整测试集上运行推理，并比较它们的输出。NumPy 提供了 allclose 函数用于此目的：

import numpy as np

for i in range(10):
    np.allclose(ort_outputs[i], outputs[i].cpu(), rtol=1e-05, atol=1e-04)

在前面的代码片段中，我们假设已经用原始模型和转换后模型在测试集上运行推理，并把输出保存到了两个数组中：outputs 和 ort_outputs。当两个数组在由 rtol，相对容差，以及 atol，绝对容差，定义的容忍范围内逐元素相等时，allclose 函数会返回 True，这些通常都是很小的值。allclose 会计算如下公式：

absolute(ort_outputs[i] - outputs[i]) <= (atol + rtol * 
➥absolute(outputs[i]))

导出并验证模型后，我们可以运行推理。首先创建一个 ONNX Runtime session：

import onnxruntime
import numpy

sess_options = onnxruntime.SessionOptions()
session = onnxruntime.InferenceSession(export_model_path, 
➥sess_options, providers=['CPUExecutionProvider'])

然后在 session 中运行推理：

inputs = tokenizer(squad["question"][0], 
➥squad["context"][0], return_tensors="np")
ort_inputs = {
    'input_ids':  full_inputs['input_ids'],
    'input_mask': full_inputs['attention_mask'],
    'segment_ids': full_inputs['token_type_ids']
}
ort_outputs = session.run(None, ort_inputs

输入会被分词为 NumPy 数组，而不是原始模型中的 PyTorch tensor。然后我们显式设置该架构期望的输入，例如这个 BERT 模型中的 input_ids、input_mask 和 segment_ids。

如果你在 notebook 中针对完整测试集运行两个模型，也就是原始模型和 ONNX 转换模型的 benchmark，会看到类似结果：

PyTorch CPU Average inference time = 455.49 ms
OnnxRuntime cpu Average inference time = 373.10 ms

ONNX 的平均推理时间大约低 18%。

最后，notebook 代码使用 ONNX API 进一步优化转换后的模型，使用 ONNX optimizer 类的 optimize_model 函数：

from onnxruntime.transformers import optimizer

optimized_model_path = os.path.join(output_dir, '
➥bert-base-uncased.onnx_opt_cpu.onnx')
optimized_model = optimizer.optimize_model(export_model_path,
➥model_type='bert', num_heads=12, hidden_size=768)
optimized_model.save_model_to_file(optimized_model_path)

在这个优化模型上运行 benchmark，会得到如下结果：

OnnxRuntime cpu Average inference time = 355.48 ms

5.6 在 GPU 上将 ONNX 用于 LLM

现在我们看看如何在 GPU 上运行 ONNX 模型。作为第一个示例，我们仍然考虑 bert-base-uncased，也就是 5.5 节中的同一个模型，但这次环境中安装了 NVIDIA GPU 和 CUDA 驱动。本示例完整代码可以在配套 Colab notebook 中找到。我们同样使用 Hugging Face Transformers 和 Datasets 库。除了 onnx 库之外，你还需要面向 GPU 的 onnxruntime：

pip install onnx onnxruntime-gpu

从 HF Hub 下载模型的代码几乎与 4.5 节示例相同：

import torch
from transformers import AutoModelForQuestionAnswering, BertTokenizer

device = torch.device("cuda")
model_id = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_id)
model = AutoModelForQuestionAnswering.from_pretrained(model_id)
model.eval()
model.to(device)

这里唯一的区别是，我们把模型加载到 GPU 内存中，因为我们也会使用硬件加速将模型转换为 ONNX。

和 5.5 节一样，我们使用 SQuAD 作为测试和验证的参考数据集。关于下载样本子集的代码，请参考上一节。

转换代码与 CPU 情况相同，因此这里不再重复；请参考 5.5 节或本节的配套 notebook。使用 ONNX 转换后模型进行推理也类似，但必须将 ONNX inference session 配置为使用 CUDA execution provider，如下所示：

session = onnxruntime.InferenceSession(export_model_path, sess_options,
➥providers=['CUDAExecutionProvider'])

这就是以标准方式在 GPU runtime 上运行 ONNX 转换模型所需的一切。还有其他需要考虑的事项，但在探索它们之前，我们先看另一个示例。

5.6.1 在 GPU 上将 ONNX 用于 GPT

前一个示例和 5.5 节中的示例使用的是 encoder-only 模型。现在，我们来看一个 decoder-only 模型：GPT-2 small。本节代码包含在配套 Colab notebook 中。

先从 Hugging Face Hub 下载模型：

import torch
from transformers import GPT2Tokenizer, AutoModelForCausalLM

model_id = 'openai-community/gpt2'
tokenizer = GPT2Tokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
device = torch.device("cuda")
model.eval().to(device)

虽然这不是必需的，但我们把模型 checkpoint 移动到了 GPU 内存，并将使用 GPU 将模型转换为 ONNX。你也可以在 CPU 上运行转换。现在创建一个目录，用来存放转换后的模型：

import os

output_dir = os.path.join(".", "onnx_models")
onnx_model_filename = 'bert-base-uncased.onnx'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
export_model_path = os.path.join(output_dir, onnx_model_filename)

接下来，准备一个样本输入。对于这个用例，我们将 tokenizer 的 return_attention_mask 参数设置为 False，因为 attention mask 不会被使用：

tokenized_inputs = tokenizer("The story so far: in the beginning, 
➥the universe was created.",
                             return_attention_mask=False,
                             return_tensors="pt")
tokenized_inputs.to(device)
inputs_sample = {
        'input_ids':  tokenized_inputs['input_ids']
    }

最后，开始转换流程：

with torch.no_grad():
  torch.onnx.export(model,
                    inputs_sample,
                    export_model_path,
                    export_params=True,
                    opset_version=15,
                    input_names=['input_ids']
                    )

我们可以像前面的 BERT 示例一样，使用导出的模型运行推理：创建 inference session，选择 CUDA runtime provider，然后调用 run 方法：

import onnxruntime
import numpy

session = onnxruntime.InferenceSession(export_model_path,
➥providers=["CUDAExecutionProvider"])
onnx_input_ids = tokenizer("The story so far: in the
➥beginning, the universe was created.", 
                           return_attention_mask=False, 
                           return_tensors="np")
ort_inputs = {
    "input_ids": onnx_input_ids['input_ids']
}

ort_outputs = session.run(None, ort_inputs)

到目前为止的示例中，我们都假设把模型转换为 ONNX 会自动应用最佳可能优化。这通常是对的，但有时候需要额外步骤。你通常可以通过程序完成这些步骤；手动检查生成的 ONNX 图应当作为最后手段。其中一种方式，不是唯一方式，后面会介绍更多，是使用 ONNX Runtime 库内置的 Transformer optimizer。继续 GPT-2 示例，下面展示如何将它应用到已经转换好的模型上：

from onnxruntime.transformers import optimizer

optimized_model_path = os.path.join(output_dir, 'gpt-2-onnx_opt_gpu.onnx')
optimized_model = optimizer.optimize_model(export_model_path, 
                                           model_type='gpt2', 
                                           use_gpu=True, 
                                           num_heads=12, 
                                           hidden_size=768)
optimized_model.save_model_to_file(optimized_model_path)

我们和前面的 bert-base-uncased 示例一样，使用 optimizer 的 optimize_model 方法。我们传入转换后模型的完整路径、LLM 家族类型、head 数量和 hidden size，并指定 GPU 优化。你也可以设置图优化级别：0，表示无优化，也是默认值；1，表示 basic；2，表示 extended；99，表示 all，包括 layout。相同代码也适用于 GPT-2 之外其他架构的 LLM。Basic optimization 会执行 constant folding，也就是静态计算图中只依赖常量 initializer 的部分；删除冗余节点，例如 identity、slice、unsqueeze 和 dropout；并应用保持语义的融合。Extended optimization 会在图分区之后应用更复杂的融合，并且只作用于分配给 CPU、CUDA 或 ROCm，也就是 AMD Radeon Open Compute Ecosystem，execution provider 的节点。Layout optimization 会对适用节点改变数据布局以获得更高性能；它在图分区之后运行，并且只作用于分配给 CPU execution provider 的节点。

ONNX 优化可以离线完成，也可以在线完成。前面的代码片段就是离线优化示例。另一种方式是在推理时通过设置 ONNX Runtime session mode 启用相同优化，如下所示：

import onnxruntime as rt

sess_options = rt.SessionOptions()

sess_options.graph_optimization_level =
➥rt.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
sess_options.optimized_model_filepath = "<model_path\optimized_model.onnx>"

session = rt.InferenceSession(None, sess_options)

优化级别由 GraphOptimizationLevel 类定义：

GraphOptimizationLevel.ORT_DISABLE_
GraphOptimizationLevel.ORT_ENABLE_BASIC 
GraphOptimizationLevel.ORT_ENABLE_EXTENDED 
GraphOptimizationLevel.ORT_ENABLE_ALL

在在线模式下，ONNX Runtime 会在 session 初始化时、推理运行之前应用所有已启用的图优化。这会增加模型启动时间，尤其是复杂模型，因此可能不适合生产环境。在离线模式下，ONNX Runtime 会应用优化，然后把优化后的模型序列化到磁盘。后续启动会更快，因为模型已经优化完成。

5.6.2 I/O Binding

在继续之前，考虑一下使用 execution provider，也就是 CPU 和 GPU，运行 ONNX 模型时的输入和输出数据。到目前为止，你已经知道 ONNX 模型由一个计算图构成，该计算图由针对不同硬件优化的 operator kernel 实现；同时 ONNX Runtime，也就是 ORT，会通过面向特定目标的 execution provider，CPU、GPU 等，来编排执行。但被消费和产生的数据怎么办？

默认情况下，无论数据结构是什么，NumPy array、dictionary，还是 NumPy array 的 list，ORT 都会把数据放在 CPU 上。在没有硬件加速时，这没有问题；但当 GPU 可用，并且 inference session 被配置为优先在 GPU 上运行时，这并不理想。在这种情况下，数据会在 CPU 和 GPU 之间拷贝，从而增加不必要开销。ORT 提供了一种自定义数据结构，它支持 ONNX 数据格式，并允许你把数据放到某个设备上，例如支持 CUDA 的 GPU。它叫 IOBinding。借助它，输入和输出数据，或者根据你的配置只绑定其中一种，可以放在 ONNX 图运行的同一设备上。

你可以通过程序配置 IOBinding，接下来会看到。为了让概念保持简单，我们会使用一个小型 PyTorch 模型，它对两个 tensor 求和；但同样方法也适用于更复杂的模型，包括 LLM。先定义模型：

import torch

def model():
    class Model(torch.nn.Module):
        def __init__(self):
            super(Model, self).__init__()

        def forward(self, x, y):
            return x.add(y)

    return Model()  

现在通过程序将其转换为 ONNX：

MODEL_FILE = '.model.onnx'

type = torch.float32
sample_x = torch.ones(3, dtype=type)
sample_y = torch.zeros(3, dtype=type)

torch.onnx.export(model(), (sample_x, sample_y), 
➥MODEL_FILE, input_names=["x", "y"], output_names=["z"],
                    ➥dynamic_axes={"x": {0 : "array_length_x"}, "y":
                    ➥{0: "array_length_y"}})

像往常一样创建 inference session，假设目标环境有 GPU 和 CUDA：

import onnxruntime as rt

providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = rt.InferenceSession(MODEL_FILE, providers=providers)

新的部分从创建 session 之后开始。我们不再使用 run 运行推理，而是使用 run_with_iobinding。首先准备 IOBinding 所需的数据。从 inference session 创建一个 IOBinding 实例：

binding = session.io_binding() 

创建两个包含随机值的 tensor，并把它们移动到 GPU：

DEVICE = 'cuda'
x = torch.rand(5).to(DEVICE)
y = torch.rand(5).to(DEVICE)
x_tensor = x.contiguous()
y_tensor = y.contiguous()

然后告诉 ONNX Runtime 把模型输入加载到 CUDA 设备上：

DEVICE_INDEX = 0 
np_type = np.float32

binding.bind_input(
    name='x',
    device_type=DEVICE,
    device_id=DEVICE_INDEX,
    element_type=np_type,
    shape=tuple(x_tensor.shape),
    buffer_ptr=x_tensor.data_ptr(),
    )

binding.bind_input(
    name='y',
    device_type=DEVICE,
    device_id=DEVICE_INDEX,
    element_type=np_type,
    shape=tuple(y_tensor.shape),
    buffer_ptr=y_tensor.data_ptr(),
    )

可以按如下方式使用 IOBinding 运行推理：

session.run_with_iobinding(binding)

在这个示例中，只有输入数据被加载到了设备上；输出仍然会被移动到 CPU。如果也想对输出使用 IOBinding，需要在运行推理之前为模型输出分配一个 PyTorch tensor：

z_tensor = torch.empty(x_tensor.shape, dtype=torch_type,
➥device=DEVICE).contiguous()

然后像绑定输入 tensor 一样绑定它：

binding.bind_output(
    name='z',
    device_type=DEVICE,
    device_id=DEVICE_INDEX,
    element_type=np_type,
    shape=tuple(z_tensor.shape),
    buffer_ptr=z_tensor.data_ptr(),
)

第 7 章和第 8 章会展示更多包含更复杂数据的示例。

总结

• ONNX 是一种开放标准格式，用于在不同框架和硬件平台之间实现机器学习模型互操作。
• ONNX Runtime，也就是 ORT，可以在 CPU、GPU 和边缘设备上高效运行 ONNX 模型。
• ONNX 算子是构建块，可以描述来自大多数框架的机器学习模型。
• ONNX 支持多种数据类型，包括 float32、int8、int16、int32、int64，以及 bfloat16 等专用格式。
• ONNX Runtime execution provider 通过使用 NVIDIA CUDA 等硬件加速器来优化性能。
• 你可以在转换过程中离线优化 ONNX 模型，也可以在初始化 inference session 时在线优化。
• ONNX 图优化范围从 constant folding 到 node fusion 和 layout transformation。
• ONNX Runtime 中的 IOBinding 通过让数据保留在目标执行设备上，消除 CPU–GPU 拷贝。
• ONNX 转换后的模型通常比原始框架实现获得更快推理时间。
• 验证 ONNX 模型意味着验证 schema 一致性，并确保输出与原始模型匹配。
• ONNX Runtime 支持多种语言 binding，但 Python 提供最完整的 API。
• 你可以使用 onnxruntime.transformers.optimizer 模块优化 ONNX Transformer 模型。
• 将语言模型转换为 ONNX 可以保持功能，并支持跨多样化生产环境部署。