在 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.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 檔案安裝方法。或者，可以按照這些說明使用 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 上更新