本文将讨论 ONNX 运行时，它是加速稳定扩散推理的最有效方法之一。在 A100 GPU 上，运行 SDXL 30 个去噪步骤生成 1024 x 1024 图像的速度可快至 2 秒。然而，ONNX 运行时依赖于多个活动部件，在不断发展的生态系统中，安装其所有依赖项的正确版本可能非常棘手。在这里，我将与大家分享我在调试过程中遇到的困难，希望能为你节省时间。虽然具体的版本和命令可能很快就会过时，但高层次的概念应该会在较长时间内保持相关。

什么是 ONNX？

ONNX 实际上指的是 ML 堆栈中两个不同（但相关）的部分：

1. ONNX 是一种用于存储机器学习模型的格式。它是开放神经网络交换（Open Neural Network Exchange）的缩写，顾名思义，其主要目标是实现跨平台互操作。ONNX 是一种自包含格式：它同时存储模型权重和架构。这意味着一个 .onnx 文件就包含了运行推理所需的全部信息。无需编写任何额外代码来定义或加载模型，只需将其传递给运行时即可。
2. ONNX 也是运行 ONNX 格式模型的运行时。它可以运行模型。你可以把它看作是架构无关的 ONNX 格式与运行推理的实际硬件之间的中介。每种支持的加速器类型都有一个独立的运行时版本（参见此处的完整列表）。不过请注意，ONNX 运行时并不是使用 ONNX 格式模型运行推理的唯一方法，它只是其中一种方法。制造商可以选择构建自己的运行时，并针对自己的硬件进行超级优化。例如，英伟达公司的 TensorRT 就是 ONNX 运行时的替代方案。

本文主要介绍使用 ONNX 运行时运行稳定扩散模型。虽然高级概念可能是永恒的，但请注意，ML 工具生态系统在不断变化，因此确切的工作流程或代码片段可能会过时（本文写于 2024 年 5 月）。我将重点介绍 Python 的实现，但要注意的是，ONNX 运行时也可以在 C++、C#、Java 或 JavaScript 等其他语言中运行。

ONNX 运行时的优点

• 推理速度与互操作性之间的平衡。虽然 ONNX 运行时并不总是所有类型硬件的最快解决方案，但对于大多数类型的硬件来说，它已经足够快了。如果你在异构机群上运行模型，而又没有资源针对每种不同的加速器进行微优化，这一点就特别有吸引力。
• 广泛采用，作者可靠。ONNX 由微软开源，目前微软仍在对其进行维护。它被广泛采用，并很好地融入了更广泛的 ML 生态系统。例如，Hugging Face 的 Optimum 库允许你定义和运行 ONNX 模型管道，其语法让人联想到其流行的 transformers 和 diffusers 库。

ONNX 运行时的缺点

• 工程开销。与直接在 PyTorch 中运行推理相比，ONNX 运行时需要将模型编译成 ONNX 格式（稳定扩散模型可能需要 20-30 分钟）并安装运行时。
• 操作集受限。ONNX 格式并不支持所有 PyTorch 操作（它甚至比 TorchScript 更严格）。如果你的模型使用了不支持的操作，你就必须重新实现相关部分，或者完全放弃 ONNX。
• 脆弱的安装和设置。由于 ONNX 运行时会将 ONNX 格式转换为特定架构的指令，因此要获得正确的软件版本组合使其正常运行可能会很棘手。例如，如果在英伟达™（NVIDIA®）GPU 上运行，则需要确保以下方面的兼容性：（1）操作系统；（2）CUDA 版本；（3）cuDNN 版本；（4）ONNX 运行时版本。有一些有用的资源，如 CUDA 兼容性矩阵，但你可能最终还是要浪费几个小时来寻找在特定时间点上有效的神奇组合。
• 硬件限制。虽然 ONNX 运行时可以在许多架构上运行，但它不能像纯 PyTorch 模型那样在所有架构上运行。例如，目前（2024 年 5 月）还不支持Google Cloud TPU 或 AWS Inferentia 芯片。

乍一看，缺点比优点多，但不要气馁--如后面所示，模型延迟的改善是显著的，也是值得的。

如何安装 ONNX 运行时

方案 1：从源代码安装

如上所述，ONNX 运行时需要与许多软件兼容。如果你想走在最前沿，获得最新版本的最佳方法就是按照官方 Github 代码库中的说明进行安装。特别是对于 Stable Diffusion，该文件夹包含安装说明和生成镜像的示例脚本。从源代码开始构建需要相当长的时间（约 30 分钟）。

此解决方案可在亚马逊 EC2 实例（g5.2xlarge，配备 A10G GPU）上无缝运行。它通过使用自带正确依赖项的 Docker 镜像，避免了下文讨论的兼容性问题。

方案 2：通过 PyPI 安装

在生产过程中，你很可能需要从 PyPI 安装稳定版本的 ONNX 运行时，而不是从源代码安装最新版本。特别是对于 Python，有两个不同的库（一个用于 CPU，一个用于 GPU）。下面是为 CPU 安装的命令：

pip install onnxruntime

下面是为 GPU 安装它的命令：

pip install onnxruntime-gpu

千万不要同时安装这两种软件。同时安装这两种软件可能会导致错误信息或行为，而这些错误信息或行为并不容易追溯到根本原因。ONNX 运行时可能只是无法识别 GPU 的存在，而这在 onnxruntime-gpu 确实已安装的情况下会显得令人惊讶。

解决兼容性问题

在理想情况下，pip install onnxruntime-gpu 就能解决这个问题。但实际上，机器上的其他软件（包括操作系统、特定硬件驱动程序和 Python 版本）之间有很强的兼容性要求。

假设你想使用本文撰写时的最新版 ONNX 运行时（1.17.1）。那么我们需要对准哪些星星才能实现这一目标呢？

以下是一些最常见的不兼容来源，可以帮助你设置环境。、

CUDA 兼容性

使用英伟达™（NVIDIA®）GPU，CUDA 是英伟达™（NVIDIA®）GPU 上的并行计算平台，是机器学习工作流所必需的。每个版本的 ONNX 运行时都只与特定的 CUDA 版本兼容，这一点你可以从兼容性矩阵中看到。

根据该矩阵，最新的 ONNX 运行时版本（1.17）兼容 CUDA 11.8 和 CUDA 12。但你需要注意其中的细节：默认情况下，ONNX 运行时 1.17 希望使用 CUDA 11.8。不过，现在（2024 年 5 月）的大多数虚拟机都配备了 CUDA 12.1（可通过运行 nvcc --version 查看版本）。对于这种特殊设置，你必须用以下命令替换通常的 pip install onnxruntime-gpu：

pip install onnxruntime-gpu==1.17.1 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/

请注意，与其任由机器上安装的 CUDA 版本摆布，不如在 Docker 容器中完成工作。你只需选择包含所需 Python 和 CUDA 版本的映像即可。例如：

docker run --rm -it --gpus all nvcr.io/nvidia/pytorch:23.10-py3rm -it --gpus all nvcr.io/nvidia/pytorch:23.10-py3

操作系统 + Python + pip 的兼容性

将讨论与体系结构无关的兼容性问题（即无论使用哪种目标加速器，都会遇到这些问题）。归根结底，就是要确保你的软件（操作系统、Python 安装和 pip 安装）与你所需的 ONNX 运行时库版本兼容。

pip 版本： 除非你使用的是传统代码或系统，否则最安全的做法是将 pip 升级到最新版本：

python -m pip install --upgrade pip

Python 版本： 截至 2024 年 5 月，最不可能让你头疼的 Python 版本是 3.10（大多数虚拟机的默认版本）。同样，除非你使用的是传统代码，否则肯定希望至少是 3.8（因为 3.7 已于 2023 年 6 月废弃）。

操作系统： 操作系统版本也会妨碍你安装所需的库，这一点让我很吃惊，尤其是我使用的是最标准的 EC2 实例。而且，要找出操作系统版本是罪魁祸首也并非易事。

在此，我将带大家回顾一下我的调试过程。首先，我用以下命令安装了 onnxruntime-gpu（因为我的机器上安装了 CUDA 12.1）：

pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/

从表面上看，这应该是安装 PyPI 上最新版本的库。但实际上，这会安装与你当前设置（操作系统 + Python 版本 + pip 版本）兼容的最新版本。对我来说，当时的版本是 onnxruntime-gpu===1.16.0（而 1.17.1 是最新版本）。在不知情的情况下安装了旧版本，结果只是 ONNX 运行时无法检测到 GPU，没有其他线索。在无意中发现版本比预期的要旧之后，我明确要求安装更新的版本：

pip install onnxruntime-gpu==1.17.1 --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/

结果，pip 发来一条消息，抱怨我请求的版本实际上不可用（尽管已在 PyPI 上列出）：

ERROR: Could not find a version that satisfies the requirement onnxruntime-gpu==1.17.1 (from versions: 1.12.0, 1.12.1, 1.13.1, 1.14.0, 1.14.1, 1.15.0, 1.15.1, 1.16.0, 1.16.1, 1.16.2, 1.16.3)
ERROR: No matching distribution found for onnxruntime-gpu==1.17.1

要了解为什么没有安装最新版本，你可以通过一个标志使 pip 成为冗长的： pip install ... -vvv. 这将显示 pip 循环过的所有 Python 轮子，以便找到与你的系统兼容的最新版本。下面是我的输出结果：

Skipping link: none of the wheel's tags (cp35-cp35m-manylinux1_x86_64) are compatible (run pip debug --verbose to show compatible tags): https://files.pythonhosted.org/packages/26/1a/163521e075d2e0c3effab02ba11caba362c06360913d7c989dcf9506edb9/onnxruntime_gpu-0.1.2-cp35-cp35m-manylinux1_x86_64.whl (from https://pypi.org/simple/onnxruntime-gpu/)
Skipping link: none of the wheel's tags (cp36-cp36m-manylinux1_x86_64) are compatible (run pip debug --verbose to show compatible tags): https://files.pythonhosted.org/packages/52/f2/30aaa83bc9e90e8a919c8e44e1010796eb30f3f6b42a7141ffc89aba9a8e/onnxruntime_gpu-0.1.2-cp36-cp36m-manylinux1_x86_64.whl (from https://pypi.org/simple/onnxruntime-gpu/)
Skipping link: none of the wheel's tags (cp37-cp37m-manylinux1_x86_64) are compatible (run pip debug --verbose to show compatible tags): https://files.pythonhosted.org/packages/a2/05/af0481897255798ee57a242d3989427015a11a84f2eae92934627be78cb5/onnxruntime_gpu-0.1.2-cp37-cp37m-manylinux1_x86_64.whl (from https://pypi.org/simple/onnxruntime-gpu/)
Skipping link: none of the wheel's tags (cp35-cp35m-manylinux1_x86_64) are compatible (run pip debug --verbose to show compatible tags): https://files.pythonhosted.org/packages/17/cb/0def5a44db45c6d38d95387f20057905ce2dd4fad35c0d43ee4b1cebbb19/onnxruntime_gpu-0.1.3-cp35-cp35m-manylinux1_x86_64.whl (from https://pypi.org/simple/onnxruntime-gpu/)
Skipping link: none of the wheel's tags (cp36-cp36m-manylinux1_x86_64) are compatible (run pip debug --verbose to show compatible tags): https://files.pythonhosted.org/packages/a6/53/0e733ebd72d7dbc84e49eeece15af13ab38feb41167fb6c3e90c92f09cbb/onnxruntime_gpu-0.1.3-cp36-cp36m-manylinux1_x86_64.whl (from https://pypi.org/simple/onnxruntime-gpu/)
...

括号中列出的标记是 Python 平台兼容性标记，简而言之，每个 Python 轮都带有一个标签，表明它能在什么系统上运行。例如，cp35-cp35m-manylinux1_x86_64 需要 CPython 3.5、manylinux1 下的一组（较旧的）Linux 发行版，以及 64 位 x86 兼容处理器。

由于我想在 Linux 机器上运行 Python 3.10（因此过滤了 cp310.*manylinux.*），我只能为 onnxruntime-gpu 库选择一个可能的轮子，标记如下：

cp310-cp310-manylinux_2_28_x86_64

运行 pip debug --verbose 可以获得与系统兼容的标签列表。下面是我的部分输出结果：

cp310-cp310-manylinux_2_26_x86_64
cp310-cp310-manylinux_2_25_x86_64
cp310-cp310-manylinux_2_24_x86_64
cp310-cp310-manylinux_2_23_x86_64
cp310-cp310-manylinux_2_22_x86_64
cp310-cp310-manylinux_2_21_x86_64
cp310-cp310-manylinux_2_20_x86_64
cp310-cp310-manylinux_2_19_x86_64
cp310-cp310-manylinux_2_18_x86_64
cp310-cp310-manylinux_2_17_x86_64
...

换句话说，我的操作系统太老了（它支持的最大 linux 标签是 manylinux_2_26，而 onnxruntime-gpu 库的 Python 3.10 wheel 需要 manylinux_2_28）。从 Ubuntu 20.04 升级到 Ubuntu 24.04 就解决了这个问题。

如何使用 ONNX 运行时运行稳定版 Diffusion

一旦（最终）安装了 ONNX 运行时，使用 Stable Diffusion 生成图像需要以下两个步骤：

1. 将 PyTorch 模型导出到 ONNX（这可能需要 30 分钟！）。
2. 将 ONNX 模型和输入（文本提示和其他参数）传递给 ONNX 运行时。

方案 1：使用微软官方脚本

如前所述，使用 ONNX 运行时资源库中的官方示例脚本对我来说是开箱即用。如果你按照他们的安装说明进行安装，你甚至无需处理上文提到的兼容性问题。安装完成后，生成映像就变得非常简单了：

python3 demo_txt2img_xl.py "starry night over Golden Gate Bridge by van gogh""starry night over Golden Gate Bridge by van gogh"

在底层，该脚本使用 Hugging Face 的 diffusers 库定义 SDXL 模型，将其导出为 ONNX 格式（可能需要 30 分钟！），然后调用 ONNX 运行时。

方案 2：使用 Hugging Face 的 Optimum 库

Optimum 库提供了很多便利，允许你在各种加速器上运行模型，同时使用著名的转换器和扩散器库中熟悉的管道 API。特别是对于 ONNX，SDXL 的推理代码就是这样的：

from optimum.onnxruntime import ORTStableDiffusionXLPipeline
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
base = ORTStableDiffusionXLPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Leonardo da Vinci"
image = base(prompt).images[0]
# Don't forget to save the ONNX model
save_directory = "sd_xl_base"
base.save_pretrained(save_directory)

不过，在实际使用中，我对 Optimum 库颇有微词。首先，安装并非易事；如果天真地按照 README 文件中的安装说明进行安装，就会遇到上文所述的不兼容问题。这本身并不是 Optimum 的错，但它确实在本已脆弱的设置上又增加了一层抽象。Optimum 安装可能会调用与你的设置冲突的 onnxruntime 版本。

即使我解决了兼容性问题，也无法使用 Optimum 的 ONNX 接口在 GPU 上运行 SDXL 推理。上面的代码片段（直接取自 Hugging Face 教程）由于形状不匹配而失败，这可能是 PyTorch → ONNX 转换中的错误造成的：

[ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Add node.while running Add node.
Name:'/down_blocks.1/attentions.0/Add'
Status Message: /down_blocks.1/attentions.0/Add: left operand cannot broadcast on dim 3 LeftShape: {2,64,4096,10}, RightShape: {2,640,64,64}

减少延迟

正如承诺的那样，为 ONNX 运行时所付出的努力是值得的。在 A100 GPU 上，推理时间从 7-8 秒（运行普通 PyTorch 时）缩短到 2 秒左右。这与 TensorRT（ONNX 的 NVIDIA 专用替代方案）相当，比 torch.compile（PyTorch 的本地 JIT 编译）快约 1 秒。

据报道，换用性能更强的 GPU（如 H100）后，使用专用运行时运行模型的收益会更高。

结论

ONNX 运行时有望显著提高延迟，但也带来了不小的工程开销。它还面临着静态编译的典型权衡问题：推理速度快了很多，但图形无法动态修改（这与 peft 等动态适配器不符）。一旦你通过了实验阶段，并准备投资于高效的生产代码，ONNX 运行时和类似的编译方法就值得添加到你的管道中。