在 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 一起用于推理和评估，例如管道。使用 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 一起使用。原因如下

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

使用 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")

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

观察到的时间增益

我们测试了三个带有解码过程的常见模型：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 模型运行的。如果您对 进一步加速 感兴趣，如果您的 GPU 具有混合精度功能，您可以使用 ORTOptimizer 优化图并将模型转换为 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 根据指定的输入形状构建其引擎。一个大问题是构建引擎可能非常耗时，特别是对于大型模型。因此，作为一种变通方法，建议使用动态形状构建 TensorRT 引擎。这可以避免为新的小尺寸和大尺寸重建引擎，这在模型部署进行推理后是不希望发生的。

为此，我们使用提供程序的选项 trt_profile_min_shapes、trt_profile_max_shapes 和 trt_profile_opt_shapes 来指定引擎的最小、最大和最佳形状。例如，对于 GPT2，我们可以使用以下形状

provider_options = {
    "trt_profile_min_shapes": "input_ids:1x1,attention_mask:1x1,position_ids:1x1",
    "trt_profile_opt_shapes": "input_ids:1x1,attention_mask:1x1,position_ids:1x1",
    "trt_profile_max_shapes": "input_ids:1x64,attention_mask:1x64,position_ids:1x64",
}

将引擎缓存路径传递到提供程序选项中，这样引擎就可以一次构建，然后完全用于推理。

例如，对于文本生成，引擎可以构建为

>>> import os
>>> 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",
...     "trt_profile_min_shapes": "input_ids:1x1,attention_mask:1x1,position_ids:1x1",
...     "trt_profile_opt_shapes": "input_ids:1x1,attention_mask:1x1,position_ids:1x1",
...     "trt_profile_max_shapes": "input_ids:1x64,attention_mask:1x64,position_ids:1x64",
... }

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

引擎存储为

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，可以按照静态量化指南中解释的方式执行静态量化。

在下面的代码示例中，执行静态量化后，结果模型将加载到 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 一起用于推理和评估，例如管道。

TensorRT 对量化模型的限制

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

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

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

观察到的时间增益

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

即将推出！

< > 在 GitHub 上更新