ONNX

ONNX 是一个开放标准，定义了一组通用的运算符和一个文件格式，用于表示不同框架（包括 PyTorch 和 TensorFlow）中的深度学习模型。当模型导出到 ONNX 时，运算符会构建一个计算图（或中间表示），该图表示数据在模型中的流动。标准化的运算符和数据类型使得在框架之间切换变得容易。

Optimum 库使用配置对象将模型导出到 ONNX，这些配置对象支持多种架构，并且可以轻松扩展。如果某个模型不受支持，请随时为 Optimum 做出贡献。

导出到 ONNX 的好处包括以下几点。

图优化和量化，以提高推理性能。
使用 ORTModel API 通过 ONNX Runtime 运行模型。
为 ONNX 模型使用优化的推理管道。

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

对于本地模型，请确保模型权重和 tokenizer 文件保存在同一目录中，例如 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 上更新