使用 Hugging Face Optimum 將 Transformers 模型轉換為 ONNX

ONNX 或開放神經網路交換 (Open Neural Network eXchange) 是一種表示機器學習模型的開放標準和格式。ONNX 定義了一套通用的運算元和一個通用的檔案格式，用於在包括 PyTorch 和 TensorFlow 在內的各種框架中表示深度學習模型。

當模型匯出為 ONNX 格式時，這些運算元被用來構建一個計算圖 (通常稱為“中間表示”)，它表示資料在神經網路中的流動過程。

重要提示： ONNX 不是一個執行時。ONNX 僅是一種表示格式，可以與像 ONNX Runtime 這樣的執行時一起使用。您可以在這裡找到支援的加速器列表。

➡️瞭解更多關於 ONNX 的資訊。

Hugging Face Optimum 是一個開源庫，是 Hugging Face Transformers 的擴充套件，它提供了一個統一的效能最佳化工具 API，以在加速硬體上實現訓練和執行模型的最大效率，包括用於在 Graphcore IPU 和 Habana Gaudi 上最佳化效能的工具包。

Optimum 可用於轉換、量化、圖最佳化、加速訓練和推理，並支援 transformers pipelines。

下面您可以看到一個典型的開發者旅程，展示瞭如何利用 Optimum 和 ONNX。

➡️ 瞭解更多關於 Optimum 的資訊

所有支援的 Transformers 架構列表可以在 Transformers 文件的 ONNX 部分中找到。以下是可轉換為 ONNX 並使用 Hugging Face Optimum 最佳化的最常用架構的摘錄：

• ALBERT
• BART
• BERT
• DistilBERT
• ELECTRA
• GPT Neo
• GPT-J
• GPT-2
• RoBERTa
• T5
• ViT
• XLM
• ...

➡️ 所有支援的架構

目前有三種方法可以將您的 Hugging Face Transformers 模型轉換為 ONNX。在本節中，您將學習如何使用所有三種方法匯出用於“文字分類”的 distilbert-base-uncased-finetuned-sst-2-english，從底層的 torch API 到最使用者友好的高層 optimum API。每種方法都會做完全相同的事情。

使用 torch.onnx 匯出 (底層)

torch.onnx 允許您透過 export 方法將模型檢查點轉換為 ONNX 圖。但您必須提供許多值，如 input_names、dynamic_axes 等。

您首先需要安裝一些依賴項

pip install transformers torch

使用 export 匯出我們的檢查點

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# load model and tokenizer
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
dummy_model_input = tokenizer("This is a sample", return_tensors="pt")

# export
torch.onnx.export(
    model, 
    tuple(dummy_model_input.values()),
    f="torch-model.onnx",  
    input_names=['input_ids', 'attention_mask'], 
    output_names=['logits'], 
    dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'}, 
                  'attention_mask': {0: 'batch_size', 1: 'sequence'}, 
                  'logits': {0: 'batch_size', 1: 'sequence'}}, 
    do_constant_folding=True, 
    opset_version=13, 
)

使用 transformers.onnx 匯出 (中層)

transformers.onnx 允許您透過利用配置物件將模型檢查點轉換為 ONNX 圖。這樣您就不必為 dynamic_axes 等提供複雜的配置。

您首先需要安裝一些依賴項

pip install transformers[onnx] torch

使用 transformers.onnx 匯出我們的檢查點。

from pathlib import Path
import transformers
from transformers.onnx import FeaturesManager
from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification

# load model and tokenizer
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
feature = "sequence-classification"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# load config
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=feature)
onnx_config = model_onnx_config(model.config)

# export
onnx_inputs, onnx_outputs = transformers.onnx.export(
        preprocessor=tokenizer,
        model=model,
        config=onnx_config,
        opset=13,
        output=Path("trfs-model.onnx")
)

使用 Optimum 匯出 (高層)

Optimum 推理包含使用 ORTModelForXxx 類將原生 Transformers 模型轉換為 ONNX 的方法。要將您的 Transformers 模型轉換為 ONNX，您只需向 from_pretrained() 方法傳遞 from_transformers=True，您的模型就會被載入並轉換為 ONNX，其底層利用了 transformers.onnx 包。

您首先需要安裝一些依賴項

pip install optimum[onnxruntime]

使用 ORTModelForSequenceClassification 匯出我們的檢查點

from optimum.onnxruntime import ORTModelForSequenceClassification

model = ORTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english",from_transformers=True)

使用 Optimum 進行轉換的最大好處是，您可以立即使用該 model 來執行預測，或將其載入到 pipeline 中。