Transformers 文档

导出到ONNX

导出到ONNX

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

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

本指南展示了如何使用 🤗 Optimum 将 🤗 Transformers 模型导出为 ONNX 格式,关于将模型导出为 TFLite 的指南,请参阅 导出到 TFLite 页面

导出到ONNX

ONNX (Open Neural Network eXchange) 是一个开放标准,定义了一组通用的操作符和通用的文件格式,用于在各种框架(包括 PyTorch 和 TensorFlow)中表示深度学习模型。当模型导出为 ONNX 格式时,这些操作符用于构建一个计算图(通常称为中间表示),该图表示数据在神经网络中的流动。

通过公开具有标准化操作符和数据类型的图,ONNX使得在框架之间切换变得容易。例如,在PyTorch中训练的模型可以导出为ONNX格式,然后在TensorFlow中导入(反之亦然)。

一旦导出为ONNX格式,模型可以:

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

有关现成配置的列表,请参阅 🤗 Optimum 文档

有两种方法可以将🤗 Transformers模型导出为ONNX格式,这里我们展示这两种方法:

  • 通过CLI使用🤗 Optimum导出。
  • 使用🤗 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/
< > Update on GitHub