ONNX

ONNX 是一個開放標準，它定義了一組通用的運算元和一個檔案格式，用於表示包括 PyTorch 和 TensorFlow 在內的不同框架中的深度學習模型。當一個模型被匯出到 ONNX 時，這些運算元會構建一個計算圖（或*中間表示*），它表示了資料在模型中的流動過程。標準化的運算元和資料型別使得在不同框架之間切換變得容易。

Optimum 庫透過配置物件將模型匯出為 ONNX，這些配置物件支援多種架構並且可以輕鬆擴充套件。如果某個模型不受支援，歡迎隨時向 Optimum 提交貢獻。

匯出到 ONNX 的好處包括：

透過圖最佳化和量化來提升推理效能。
使用 ORTModel API 來透過 ONNX Runtime 執行模型。
為 ONNX 模型使用最佳化的推理 pipeline。

使用 Optimum CLI 或 optimum.onnxruntime 模組將 Transformers 模型匯出為 ONNX。

Optimum CLI

執行以下命令以安裝 Optimum 和 exporters 模組。

pip install optimum[exporters]

請參閱使用 optimum.exporters.onnx 將模型匯出為 ONNX 指南，或使用以下命令瞭解所有可用引數。

optimum-cli export onnx --help

設定 --model 引數來從 Hub 匯出 PyTorch 或 TensorFlow 模型。

optimum-cli export onnx --model distilbert/distilbert-base-uncased-distilled-squad distilbert_base_uncased_squad_onnx/

您應該會看到指示進度的日誌，並顯示生成的 model.onnx 檔案的儲存位置。

Validating ONNX model distilbert_base_uncased_squad_onnx/model.onnx...
	-[✓] ONNX model output names match reference model (start_logits, end_logits)
	- Validating ONNX Model output "start_logits":
		-[✓] (2, 16) matches (2, 16)
		-[✓] all values close (atol: 0.0001)
	- Validating ONNX Model output "end_logits":
		-[✓] (2, 16) matches (2, 16)
		-[✓] all values close (atol: 0.0001)
The ONNX export succeeded and the exported model was saved at: distilbert_base_uncased_squad_onnx

對於本地模型，請確保模型權重和分詞器檔案儲存在同一目錄下，例如 local_path。將該目錄傳遞給 --model 引數，並使用 --task 來指定模型可以執行的任務。如果未提供 --task，則會使用不帶特定任務頭的模型架構。

optimum-cli export onnx --model local_path --task question-answering distilbert_base_uncased_squad_onnx/

model.onnx 檔案可以部署在任何支援 ONNX 的加速器上。下面的示例演示瞭如何使用 ONNX Runtime 載入和執行模型。

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

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert_base_uncased_squad_onnx")
>>> model = ORTModelForQuestionAnswering.from_pretrained("distilbert_base_uncased_squad_onnx")
>>> inputs = tokenizer("What am I using?", "Using DistilBERT with ONNX Runtime!", return_tensors="pt")
>>> outputs = model(**inputs)

optimum.onnxruntime

optimum.onnxruntime 模組支援以程式設計方式匯出 Transformers 模型。為特定任務例項化一個 ORTModel 並設定 export=True。使用 ~OptimizedModel.save_pretrained 來儲存 ONNX 模型。

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

>>> model_checkpoint = "distilbert/distilbert-base-uncased-distilled-squad"
>>> save_directory = "onnx/"

>>> ort_model = ORTModelForSequenceClassification.from_pretrained(model_checkpoint, export=True)
>>> tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

>>> ort_model.save_pretrained(save_directory)
>>> tokenizer.save_pretrained(save_directory)

< > 在 GitHub 上更新