在 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 之前运行 pip uninstall onnxruntime 卸载 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 一起用于推理和评估，例如 pipelines。当使用 Transformers pipeline 时，请注意应设置 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（verbose），以检查是否所有节点都确实放置在 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 一起使用。原因如下

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

使用 IOBinding 减少内存占用

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

为了避免速度减慢，🤗 Optimum 采用 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")

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

观察到的时间增益

我们使用解码过程测试了三个常用模型：GPT2 / T5-small / M2M100-418M，并且基准测试在通用的 Tesla T4 GPU 上运行（更多环境详细信息在本节末尾）。

以下是在 IOBinding 开启时使用 CUDAExecutionProvider 运行的一些性能结果。我们测试了从 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 用作通过 🤗 Optimum 优化的模型的执行提供程序的最简单方法是使用可用的 ONNX Runtime TensorrtExecutionProvider。

为了在本地环境中使用 🤗 Optimum 和 TensorRT，我们建议遵循 NVIDIA 安装指南

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

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

安装所需的软件包后，需要使用相应的路径设置以下环境变量，以便 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 引擎，然后使用大形状的输入构建 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（verbose），我们可以在日志中检查是否所有节点都确实放置在 TensorRT 执行提供程序上

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

然后，该模型可以与常用的 🤗 Transformers API 一起用于推理和评估，例如 pipelines。

将 TensorRT 执行提供程序与量化模型一起使用

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

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

>>> from optimum.onnxruntime import AutoQuantizationConfig

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

使用此 qconfig，可以按照静态量化指南中所述执行静态量化。

在下面的代码示例中，在执行静态量化后，生成的模型将加载到 ORTModel 类中，并使用 TensorRT 作为执行提供程序。需要禁用 ONNX Runtime 图优化，以便模型可以被 TensorRT 使用和优化，并且需要向 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 一起用于推理和评估，例如 pipelines。

量化模型的 TensorRT 限制

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

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

如果传递了 provider="TensorrtExecutionProvider" 并且模型没有严格遵循这些约束进行量化，则可能会引发各种错误，其中错误消息可能不明确。

观察到的时间增益

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

即将推出！

< > 在 GitHub 上更新