在 NVIDIA GPU 上加速推理

默认情况下，ONNX Runtime 在 CPU 设备上运行推理。但是，可以将受支持的操作放置在 NVIDIA GPU 上，同时将任何不受支持的操作保留在 CPU 上。在大多数情况下，这允许将代价高昂的操作放在 GPU 上并显著加速推理。

本指南将向您展示如何在 ONNX Runtime 支持的用于 NVIDIA GPU 的两个执行提供程序上运行推理

CUDAExecutionProvider：在支持 NVIDIA CUDA 的 GPU 上进行通用加速。
TensorrtExecutionProvider：使用 NVIDIA 的 TensorRT 推理引擎，通常提供最佳运行时性能。

由于 ONNX Runtime 的限制，无法在 CUDAExecutionProvider 上运行量化模型，并且只有具有静态量化的模型才能在 TensorrtExecutionProvider 上运行。

CUDAExecutionProvider

CUDA 安装

只要满足 CUDA 和 cuDNN 要求，就可以通过运行以下命令安装额外的依赖项

pip install optimum[onnxruntime-gpu]

为了避免 onnxruntime 和 onnxruntime-gpu 之间的冲突，请确保在安装 Optimum 之前未安装 onnxruntime 包，方法是运行 pip uninstall onnxruntime。

检查 CUDA 安装是否成功

在继续之前，请运行以下示例代码以检查安装是否成功

>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import AutoTokenizer

>>> ort_model = ORTModelForSequenceClassification.from_pretrained(
...   "philschmid/tiny-bert-sst2-distilled",
...   export=True,
...   provider="CUDAExecutionProvider",
... )

>>> tokenizer = AutoTokenizer.from_pretrained("philschmid/tiny-bert-sst2-distilled")
>>> inputs = tokenizer("expectations were low, actual enjoyment was high", return_tensors="pt", padding=True)

>>> outputs = ort_model(**inputs)
>>> assert ort_model.providers == ["CUDAExecutionProvider", "CPUExecutionProvider"]

如果这段代码顺利运行，恭喜你，安装成功！如果遇到以下错误或类似错误，

ValueError: Asked to use CUDAExecutionProvider as an ONNX Runtime execution provider, but the available execution providers are ['CPUExecutionProvider'].

那么 CUDA 或 ONNX Runtime 安装有问题。

将 CUDA 执行提供程序与浮点模型一起使用

对于非量化模型，使用非常简单。只需在 ORTModel.from_pretrained() 方法中指定 provider 参数。以下是一个示例

>>> from optimum.onnxruntime import ORTModelForSequenceClassification

>>> ort_model = ORTModelForSequenceClassification.from_pretrained(
...   "distilbert-base-uncased-finetuned-sst-2-english",
...   export=True,
...   provider="CUDAExecutionProvider",
... )

然后可以使用常见的 🤗 Transformers API 来执行模型推理和评估，例如管道。在使用 Transformers 管道时，请注意 device 参数应设置为在 GPU 上执行预处理和后处理，如下例所示

>>> from optimum.pipelines import pipeline
>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

>>> pipe = pipeline(task="text-classification", model=ort_model, tokenizer=tokenizer, device="cuda:0")
>>> result = pipe("Both the music and visual were astounding, not to mention the actors performance.")
>>> print(result)
# printing: [{'label': 'POSITIVE', 'score': 0.9997727274894714}]

此外，您可以传递会话选项 log_severity_level = 0（详细），以检查所有节点是否确实都放置在 CUDA 执行提供程序上。

>>> import onnxruntime

>>> session_options = onnxruntime.SessionOptions()
>>> session_options.log_severity_level = 0

>>> ort_model = ORTModelForSequenceClassification.from_pretrained(
...     "distilbert-base-uncased-finetuned-sst-2-english",
...     export=True,
...     provider="CUDAExecutionProvider",
...     session_options=session_options
... )

您应该会看到以下日志

2022-10-18 14:59:13.728886041 [V:onnxruntime:, session_state.cc:1193 VerifyEachN
odeIsAssignedToAnEp]  Provider: [CPUExecutionProvider]: [Gather (Gather_76), Uns
queeze (Unsqueeze_78), Gather (Gather_97), Gather (Gather_100), Concat (Concat_1
10), Unsqueeze (Unsqueeze_125), ...]
2022-10-18 14:59:13.728906431 [V:onnxruntime:, session_state.cc:1193 VerifyEachN
odeIsAssignedToAnEp]  Provider: [CUDAExecutionProvider]: [Shape (Shape_74), Slic
e (Slice_80), Gather (Gather_81), Gather (Gather_82), Add (Add_83), Shape (Shape
_95), MatMul (MatMul_101), ...]

在这个示例中，我们可以看到所有代价高昂的 MatMul 操作都放置在 CUDA 执行提供程序上。

使用 CUDA 执行提供程序与量化模型

由于 ONNX Runtime 的当前限制，无法将量化模型与 CUDAExecutionProvider 一起使用。原因如下：

当使用 🤗 最佳动态量化时，节点如 MatMulInteger，DynamicQuantizeLinear 可能插入 ONNX 图中，而 CUDA 执行提供程序无法使用这些节点。
当使用静态量化时，ONNX 计算图将包含浮点运算的矩阵乘法和卷积，以及用于模拟量化的量化 + 反量化操作。在这种情况下，虽然代价高昂的矩阵乘法和卷积将在 GPU 上运行，但它们将使用浮点运算，因为 CUDAExecutionProvider 无法使用量化 + 反量化节点将它们替换为使用整数运算的操作。

使用 IOBinding 减少内存占用

IOBinding 是一种在使用 GPU 时避免昂贵数据复制的有效方法。默认情况下，ONNX Runtime 会从 CPU 复制输入（即使张量已复制到目标设备），并假设输出在运行后也需要从 GPU 复制回 CPU。这些主机和设备之间的数据复制开销很昂贵，并且 **会导致比原生 PyTorch 更差的推理延迟**，尤其是在解码过程中。

为了避免速度变慢，🤗 最佳采用了 IOBinding 将输入复制到 GPU 上并在推理之前为输出预分配内存。在实例化 ORTModel 时，设置参数 use_io_binding 的值以选择在推理期间是否开启 IOBinding。如果选择 CUDA 作为执行提供程序，则默认情况下 use_io_binding 设置为 True。

如果您想关闭 IOBinding，

>>> from transformers import AutoTokenizer, pipeline
>>> from optimum.onnxruntime import ORTModelForSeq2SeqLM

# Load the model from the hub and export it to the ONNX format
>>> model = ORTModelForSeq2SeqLM.from_pretrained("t5-small", export=True, use_io_binding=False)
>>> tokenizer = AutoTokenizer.from_pretrained("t5-small")

# Create a pipeline
>>> onnx_translation = pipeline("translation_en_to_fr", model=model, tokenizer=tokenizer, device="cuda:0")

目前，IOBinding 支持任务定义的 ORT 模型，如果您希望我们添加对自定义模型的支持，请在 Optimum 的存储库中向我们提交问题。

观察到的时间收益

我们测试了三个具有解码过程的常见模型：GPT2 / T5-small / M2M100-418M，并且基准测试在一台多功能 Tesla T4 GPU 上运行（有关更多环境详细信息，请参见本节末尾）。

以下是一些使用 CUDAExecutionProvider 运行时的性能结果，其中 IOBinding 已开启。我们测试了从 8 到 512 的输入序列长度，并使用贪婪搜索和束搜索 (num_beam=5) 生成了输出。

GPT2	T5-small
M2M100-418M

以下是使用 ONNX Runtime 与 PyTorch 相比，在不同序列长度（32 / 128）和生成模式（贪婪搜索 / 束搜索）下节省时间的摘要。

seq32
序列长度：32

seq128
序列长度：128

环境

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.33.01    Driver Version: 440.33.01    CUDA Version: 11.3     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   28C    P8     8W /  70W |      0MiB / 15109MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

- Platform: Linux-5.4.0-1089-aws-x86_64-with-glibc2.29
- Python version: 3.8.10
- `transformers` version: 4.24.0
- `optimum` version: 1.5.0
- PyTorch version: 1.12.0+cu113

请注意，之前的实验使用的是直接从导出器导出的 **原生 ONNX** 模型。如果您有兴趣 **进一步加速**，可以使用 ORTOptimizer 优化图形，并在拥有混合精度功能的 GPU 的情况下将模型转换为 FP16。

TensorrtExecutionProvider

TensorRT 使用其自身的优化集，并且 **通常不支持来自 ORTOptimizer 的优化**。因此，我们建议在使用 TensorrtExecutionProvider 时使用原始 ONNX 模型 (参考)。

TensorRT 安装

使用 TensorRT 作为 🤗 最佳优化的模型执行提供程序的最简单方法是使用可用的 ONNX Runtime TensorrtExecutionProvider。

为了在本地环境中将 🤗 最佳与 TensorRT 一起使用，我们建议您遵循 NVIDIA 安装指南。

CUDA 工具包：https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html
cuDNN：https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html
TensorRT：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html

对于 TensorRT，我们建议使用 Tar 文件安装方法。或者，TensorRT 可能可以通过 pip 安装，方法是遵循这些说明。

安装完所需软件包后，需要设置以下环境变量，其中包含 ONNX Runtime 检测 TensorRT 安装的适当路径。

export CUDA_PATH=/usr/local/cuda
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-x.x/lib64:/path/to/TensorRT-8.x.x/lib

检查 TensorRT 安装是否成功

在继续之前，请运行以下示例代码以检查安装是否成功

>>> from optimum.onnxruntime import ORTModelForSequenceClassification
>>> from transformers import AutoTokenizer

>>> ort_model = ORTModelForSequenceClassification.from_pretrained(
...     "philschmid/tiny-bert-sst2-distilled",
...     export=True,
...     provider="TensorrtExecutionProvider",
... )

>>> tokenizer = AutoTokenizer.from_pretrained("philschmid/tiny-bert-sst2-distilled")
>>> inp = tokenizer("expectations were low, actual enjoyment was high", return_tensors="pt", padding=True)

>>> result = ort_model(**inp)
>>> assert ort_model.providers == ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]

如果这段代码顺利运行，恭喜您，安装成功！

如果上面的 assert 失败，或者您遇到以下警告，

Failed to create TensorrtExecutionProvider. Please reference https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html#requirements to ensure all dependencies are met.

TensorRT 或 ONNX Runtime 安装有问题。

TensorRT 引擎构建和预热

TensorRT 要求在推理之前构建其推理引擎，这需要一些时间，因为要进行模型优化和节点融合。为了避免每次加载模型时都重建引擎，ONNX Runtime 提供了一对用于保存引擎的选项：trt_engine_cache_enable 和 trt_engine_cache_path。

我们建议在使用 TensorRT 执行提供程序时设置这两个提供程序选项。使用方法如下，其中 optimum/gpt2 是使用 Optimum ONNX 导出器从 PyTorch 转换的 ONNX 模型。

>>> from optimum.onnxruntime import ORTModelForCausalLM

>>> provider_options = {
...     "trt_engine_cache_enable": True,
...     "trt_engine_cache_path": "tmp/trt_cache_gpt2_example"
... }

# the TensorRT engine is not built here, it will be when doing inference
>>> ort_model = ORTModelForCausalLM.from_pretrained(
...     "optimum/gpt2",
...     use_cache=False,
...     provider="TensorrtExecutionProvider",
...     provider_options=provider_options
... )

TensorRT 根据指定的输入形状构建其引擎。不幸的是，在当前 ONNX Runtime 实现（参考：1，2）中，每当输入形状小于之前遇到的最小形状时，引擎就会重建，反之亦然，如果输入形状大于之前遇到的最大形状，引擎也会重建。例如，如果模型将 (batch_size, input_ids) 作为输入，并且模型依次接受以下输入

input.shape: (4, 5) --> 构建引擎（第一个输入）
input.shape: (4, 10) --> 重建引擎（10 大于 5）
input.shape: (4, 7) --> 不重建（5 <= 7 <= 10）
input.shape: (4, 12) --> 重建引擎（10 <= 12）
input.shape: (4, 3) --> 重建引擎（3 <= 5）

一个主要问题是构建引擎可能很耗时，尤其是对于大型模型。因此，作为一种变通方法，建议是首先使用小型形状的输入构建 TensorRT 引擎，然后使用大型形状的输入构建引擎，以获得适用于所有中间形状的引擎。这可以避免在部署模型进行推理后，为新的小型和大型形状重建引擎，这是不必要的。

在提供程序选项中传递引擎缓存路径后，可以为所有形状构建一次引擎，并在之后的推理中完全使用它。

例如，对于文本生成，可以使用以下方法构建引擎

>>> import os
>>> from transformers import AutoTokenizer
>>> from optimum.onnxruntime import ORTModelForCausalLM

>>> os.makedirs("tmp/trt_cache_gpt2_example", exist_ok=True)
>>> provider_options = {
...     "trt_engine_cache_enable": True,
...     "trt_engine_cache_path": "tmp/trt_cache_gpt2_example"
... }

>>> ort_model = ORTModelForCausalLM.from_pretrained(
...     "optimum/gpt2",
...     use_cache=False,
...     provider="TensorrtExecutionProvider",
...     provider_options=provider_options,
... )
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/gpt2")

>>> print("Building engine for a short sequence...")
>>> text = ["short"]
>>> encoded_input = tokenizer(text, return_tensors="pt").to("cuda")
>>> output = ort_model(**encoded_input)

>>> print("Building engine for a long sequence...")
>>> text = [" a very long input just for demo purpose, this is very long" * 10]
>>> encoded_input = tokenizer(text, return_tensors="pt").to("cuda")
>>> output = ort_model(**encoded_input)

引擎将存储为

TensorRT engine cache folder

构建引擎后，可以重新加载缓存，生成过程不需要重建引擎

>>> from transformers import AutoTokenizer
>>> from optimum.onnxruntime import ORTModelForCausalLM

>>> provider_options = {
...     "trt_engine_cache_enable": True,
...     "trt_engine_cache_path": "tmp/trt_cache_gpt2_example"
... }

>>> ort_model = ORTModelForCausalLM.from_pretrained(
...     "optimum/gpt2",
...     use_cache=False,
...     provider="TensorrtExecutionProvider",
...     provider_options=provider_options,
... )
>>> tokenizer = AutoTokenizer.from_pretrained("optimum/gpt2")

>>> text = ["Replace me by any text you'd like."]
>>> encoded_input = tokenizer(text, return_tensors="pt").to("cuda")

>>> for i in range(3):
...     output = ort_model.generate(**encoded_input)
...     print(tokenizer.decode(output[0]))  # doctest: +IGNORE_RESULT

预热

构建引擎后，建议在进行推理之前执行一次或几次预热步骤，因为第一次推理运行会有一些开销。

使用 TensorRT 执行提供程序处理浮点模型

对于非量化的模型，使用起来很简单，只需在 ORTModel.from_pretrained() 中使用 provider 参数即可。例如

>>> from optimum.onnxruntime import ORTModelForSequenceClassification

>>> ort_model = ORTModelForSequenceClassification.from_pretrained(
...     "distilbert-base-uncased-finetuned-sst-2-english",
...     export=True,
...     provider="TensorrtExecutionProvider",
... )

与之前使用 CUDAExecutionProvider 相似，通过传递会话选项 log_severity_level = 0（详细），可以在日志中检查所有节点是否确实都放置在 TensorRT 执行提供程序中

2022-09-22 14:12:48.371513741 [V:onnxruntime:, session_state.cc:1188 VerifyEachNodeIsAssignedToAnEp] All nodes have been placed on [TensorrtExecutionProvider]

然后可以使用常见的 🤗 Transformers API 进行推理和评估，例如管道。

使用 TensorRT 执行提供程序处理量化模型

当涉及到量化模型时，TensorRT 仅支持使用静态量化的模型，并对权重和激活进行对称量化。

🤗 Optimum 提供了一个量化配置，可以使用 ORTQuantizer 与 TensorRT 量化的约束一起使用

>>> from optimum.onnxruntime import AutoQuantizationConfig

>>> qconfig = AutoQuantizationConfig.tensorrt(per_channel=False)

使用此 qconfig，可以按照静态量化指南中说明的步骤进行静态量化。

在下面的代码示例中，在执行静态量化后，使用 TensorRT 作为执行提供程序，将结果模型加载到 ORTModel 类中。为了让 TensorRT 使用和优化模型，需要禁用 ONNX Runtime 图优化，并且需要向 TensorRT 指定使用了 INT8 操作。

>>> import onnxruntime
>>> from transformers import AutoTokenizer
>>> from optimum.onnxruntime import ORTModelForSequenceClassification

>>> session_options = onnxruntime.SessionOptions()
>>> session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL

>>> tokenizer = AutoTokenizer.from_pretrained("fxmarty/distilbert-base-uncased-sst2-onnx-int8-for-tensorrt")
>>> ort_model = ORTModelForSequenceClassification.from_pretrained(
...     "fxmarty/distilbert-base-uncased-sst2-onnx-int8-for-tensorrt",
...     provider="TensorrtExecutionProvider",
...     session_options=session_options,
...     provider_options={"trt_int8_enable": True},
>>> )

>>> inp = tokenizer("TensorRT is a bit painful to use, but at the end of day it runs smoothly and blazingly fast!", return_tensors="np")

>>> res = ort_model(**inp)

>>> print(res)
>>> print(ort_model.config.id2label[res.logits[0].argmax()])
>>> # SequenceClassifierOutput(loss=None, logits=array([[-0.545066 ,  0.5609764]], dtype=float32), hidden_states=None, attentions=None)
>>> # POSITIVE

然后可以使用常见的 🤗 Transformers API 进行推理和评估，例如管道。

TensorRT 对量化模型的限制

如上一节所述，TensorRT 仅支持有限范围的量化模型

仅支持静态量化
权重和激活的量化范围是对称的
权重需要以 float32 格式存储在 ONNX 模型中，因此量化不会节省存储空间。实际上，TensorRT 需要插入完整的量化 + 反量化对。通常情况下，权重将以定点 8 位格式存储，并且仅对权重应用 DequantizeLinear 操作。

如果传递了 provider="TensorrtExecutionProvider" 且模型没有严格按照这些约束进行量化，可能会引发各种错误，而错误消息可能不清楚。

观察到的时间收益

可以使用 Nvidia Nsight Systems 工具来分析 GPU 上的执行时间。在分析或测量延迟/吞吐量之前，最好执行一些预热步骤。

敬请期待！

< > 在 GitHub 上更新