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 上更新