Hugging Face 徽标

Transformers 文档

导出到 ONNX

Hugging Face's logo
加入 Hugging Face 社区

并获得增强版文档体验

开始使用

导出到 ONNX

在生产环境中部署 🤗 Transformers 模型通常需要（或可以从）将模型导出到可以加载并在专用运行时和硬件上执行的序列化格式。

🤗 Optimum 是 Transformers 的扩展，它通过其 exporters 模块，支持将模型从 PyTorch 或 TensorFlow 导出到 ONNX 和 TFLite 等序列化格式。🤗 Optimum 还提供了一组性能优化工具，可在目标硬件上以最高效率训练和运行模型。

本指南演示了如何使用 🤗 Optimum 将 🤗 Transformers 模型导出到 ONNX，有关将模型导出到 TFLite 的指南，请参阅 导出到 TFLite 页面。

导出到 ONNX

ONNX（开放神经网络交换） 是一种开放标准，定义了一组通用的运算符和一个通用的文件格式，用于在各种框架中表示深度学习模型，包括 PyTorch 和 TensorFlow。当模型导出到 ONNX 格式时，这些运算符用于构建一个计算图（通常称为中间表示），它表示数据在神经网络中的流动。

通过公开具有标准化运算符和数据类型的图，ONNX 使得在框架之间轻松切换成为可能。例如，在 PyTorch 中训练的模型可以导出到 ONNX 格式，然后导入到 TensorFlow 中（反之亦然）。

导出到 ONNX 格式后，模型可以

  • 通过诸如 图优化 和 量化 等技术，针对推理进行优化。
  • 使用 ONNX Runtime 运行，通过 ORTModelForXXX 类，它遵循与您在 🤗 Transformers 中习惯使用的 AutoModel API 相同的 API。
  • 使用 优化的推理管道 运行，它与 🤗 Transformers 中的 pipeline() 函数具有相同的 API。

🤗 Optimum 通过利用配置对象来提供对 ONNX 导出的支持。这些配置对象为许多模型架构预先准备好了，并且旨在轻松扩展到其他架构。

有关预先准备好的配置列表，请参阅 🤗 Optimum 文档。

有两种方法可以将 🤗 Transformers 模型导出到 ONNX，这里我们将展示这两种方法

  • 使用 🤗 Optimum 通过 CLI 导出。
  • 使用 🤗 Optimum 以及 optimum.onnxruntime 导出。

使用 CLI 将 🤗 Transformers 模型导出到 ONNX

要将 🤗 Transformers 模型导出到 ONNX，首先安装额外的依赖项

pip install optimum[exporters]

要查看所有可用参数，请参阅 🤗 Optimum 文档，或者在命令行中查看帮助

optimum-cli export onnx --help

要将模型的检查点从 🤗 Hub 导出，例如 distilbert/distilbert-base-uncased-distilled-squad，请运行以下命令

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

上面的示例说明了如何导出 🤗 Hub 中的检查点。导出本地模型时，首先确保您将模型的权重和标记器文件都保存在同一个目录 (local_path) 中。使用 CLI 时，将 local_path 传递给 model 参数，而不是 🤗 Hub 上的检查点名称，并提供 --task 参数。您可以在 🤗 Optimum 文档 中查看支持的任务列表。如果未提供 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)

对于 🤗 Hub 上的 TensorFlow 检查点，过程相同。例如，以下是如何导出 Keras 组织 中的纯 TensorFlow 检查点

optimum-cli export onnx --model keras-io/transformers-qa distilbert_base_cased_squad_onnx/

使用 optimum.onnxruntime 将 🤗 Transformers 模型导出到 ONNX

作为 CLI 的替代方案，您可以通过编程方式将 🤗 Transformers 模型导出到 ONNX，如下所示

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

>>> model_checkpoint = "distilbert_base_uncased_squad"
>>> save_directory = "onnx/"

>>> # Load a model from transformers and export it to ONNX
>>> ort_model = ORTModelForSequenceClassification.from_pretrained(model_checkpoint, export=True)
>>> tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

>>> # Save the onnx model and tokenizer
>>> ort_model.save_pretrained(save_directory)
>>> tokenizer.save_pretrained(save_directory)

导出不支持架构的模型

如果您希望通过添加对当前无法导出的模型的支持来做出贡献，您应该首先检查它是否在 optimum.exporters.onnx 中受支持，如果不受支持，则 直接贡献到 🤗 Optimum。

使用 transformers.onnx 导出模型

transformers.onnx 现在不再维护，请使用上面介绍的 🤗 Optimum 导出模型。本节将在未来的版本中删除。

要使用 transformers.onnx 将 🤗 Transformers 模型导出到 ONNX，请安装额外的依赖项

pip install transformers[onnx]

使用 transformers.onnx 包作为 Python 模块，使用预先准备好的配置来导出检查点

python -m transformers.onnx --model=distilbert/distilbert-base-uncased onnx/

这会导出由 --model 参数定义的检查点的 ONNX 图。传递 🤗 Hub 上的任何检查点或本地存储的检查点。生成的 model.onnx 文件可以在支持 ONNX 标准的许多加速器上运行。例如，使用 ONNX Runtime 加载并运行模型，如下所示

>>> from transformers import AutoTokenizer
>>> from onnxruntime import InferenceSession

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
>>> session = InferenceSession("onnx/model.onnx")
>>> # ONNX Runtime expects NumPy arrays as input
>>> inputs = tokenizer("Using DistilBERT with ONNX Runtime!", return_tensors="np")
>>> outputs = session.run(output_names=["last_hidden_state"], input_feed=dict(inputs))

所需输出名称（如 ["last_hidden_state"]）可以通过查看每个模型的 ONNX 配置来获得。例如，对于 DistilBERT，我们有

>>> from transformers.models.distilbert import DistilBertConfig, DistilBertOnnxConfig

>>> config = DistilBertConfig()
>>> onnx_config = DistilBertOnnxConfig(config)
>>> print(list(onnx_config.outputs.keys()))
["last_hidden_state"]

对于 🤗 Hub 上的 TensorFlow 检查点，过程相同。例如，导出纯 TensorFlow 检查点，如下所示

python -m transformers.onnx --model=keras-io/transformers-qa onnx/

要导出本地存储的模型，请将模型的权重和标记器文件保存在同一个目录（例如 local-pt-checkpoint）中，然后通过将 transformers.onnx 包的 --model 参数指向所需目录来将其导出到 ONNX

python -m transformers.onnx --model=local-pt-checkpoint onnx/
< > 在 GitHub 上更新

←在 Amazon SageMaker 上运行训练 导出到 TFLite→