使用 optimum.exporters.onnx 将模型导出为 ONNX

总结

将模型导出为 ONNX 非常简单，只需

optimum-cli export onnx --model gpt2 gpt2_onnx/

查看帮助以获取更多选项

optimum-cli export onnx --help

为什么使用 ONNX？

如果您需要在生产环境中部署 🤗 Transformers 或 🤗 Diffusers 模型，我们建议将它们导出为序列化格式，以便可以在专用运行时和硬件上加载和执行。在本指南中，我们将向您展示如何将这些模型导出为 ONNX (开放神经网络交换)。

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

通过公开具有标准化运算符和数据类型的图，ONNX 可以轻松地在框架之间切换。例如，在 PyTorch 中训练的模型可以导出为 ONNX 格式，然后导入到 TensorRT 或 OpenVINO 中。

一旦导出，模型可以通过图优化和量化等技术进行推理优化。查看 optimum.onnxruntime 子包以优化和运行 ONNX 模型！

🤗 Optimum 通过利用配置对象为 ONNX 导出提供支持。这些配置对象已为多种模型架构准备就绪，并且设计为易于扩展到其他架构。

要查看支持的架构，请转到配置参考页面。

使用 CLI 将模型导出为 ONNX

要将 🤗 Transformers 或 🤗 Diffusers 模型导出为 ONNX，您首先需要安装一些额外的依赖项

pip install optimum[exporters]

Optimum ONNX 导出可以通过 Optimum 命令行使用

optimum-cli export onnx --help

usage: optimum-cli <command> [<args>] export onnx [-h] -m MODEL [--task TASK] [--monolith] [--device DEVICE] [--opset OPSET] [--atol ATOL]
                                                  [--framework {pt,tf}] [--pad_token_id PAD_TOKEN_ID] [--cache_dir CACHE_DIR] [--trust-remote-code]
                                                  [--no-post-process] [--optimize {O1,O2,O3,O4}] [--batch_size BATCH_SIZE]
                                                  [--sequence_length SEQUENCE_LENGTH] [--num_choices NUM_CHOICES] [--width WIDTH] [--height HEIGHT]
                                                  [--num_channels NUM_CHANNELS] [--feature_size FEATURE_SIZE] [--nb_max_frames NB_MAX_FRAMES]
                                                  [--audio_sequence_length AUDIO_SEQUENCE_LENGTH]
                                                  output

optional arguments:
  -h, --help            show this help message and exit

Required arguments:
  -m MODEL, --model MODEL
                        Model ID on huggingface.co or path on disk to load model from.
  output                Path indicating the directory where to store generated ONNX model.

Optional arguments:
  --task TASK           The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among: ['default', 'fill-mask', 'text-generation', 'text2text-generation', 'text-classification', 'token-classification', 'multiple-choice', 'object-detection', 'question-answering', 'image-classification', 'image-segmentation', 'masked-im', 'semantic-segmentation', 'automatic-speech-recognition', 'audio-classification', 'audio-frame-classification', 'automatic-speech-recognition', 'audio-xvector', 'image-to-text', 'zero-shot-object-detection', 'image-to-image', 'inpainting', 'text-to-image']. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder.
  --monolith            Force to export the model as a single ONNX file. By default, the ONNX exporter may break the model in several ONNX files, for example for encoder-decoder models where the encoder should be run only once while the decoder is looped over.
  --device DEVICE       The device to use to do the export. Defaults to "cpu".
  --opset OPSET         If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used.
  --atol ATOL           If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol for the model will be used.
  --framework {pt,tf}   The framework to use for the ONNX export. If not provided, will attempt to use the local checkpoint's original framework or what is available in the environment.
  --pad_token_id PAD_TOKEN_ID
                        This is needed by some models, for some tasks. If not provided, will attempt to use the tokenizer to guess it.
  --cache_dir CACHE_DIR
                        Path indicating where to store cache.
  --trust-remote-code   Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository.
  --no-post-process     Allows to disable any post-processing done by default on the exported ONNX models. For example, the merging of decoder and decoder-with-past models into a single ONNX model file to reduce memory usage.
  --optimize {O1,O2,O3,O4}
                        Allows to run ONNX Runtime optimizations directly during the export. Some of these optimizations are specific to ONNX Runtime, and the resulting ONNX will not be usable with other runtime as OpenVINO or TensorRT. Possible options:
                            - O1: Basic general optimizations
                            - O2: Basic and extended general optimizations, transformers-specific fusions
                            - O3: Same as O2 with GELU approximation
                            - O4: Same as O3 with mixed precision (fp16, GPU-only, requires `--device cuda`)

导出检查点可以按如下方式完成

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

您应该看到以下日志（以及可能来自 PyTorch/TensorFlow 的日志，这些日志为了清晰起见在此处隐藏）

Automatic task detection to question-answering.
Framework not specified. Using pt to export the model.
Using framework PyTorch: 1.12.1

Validating ONNX model...
        -[✓] 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)
All good, model saved at: distilbert_base_uncased_squad_onnx/model.onnx

这将导出由 --model 参数定义的检查点的 ONNX 图。如您所见，任务已自动检测到。这是可能的，因为模型在 Hub 上。

对于本地模型，需要提供 --task 参数，否则它将默认为不带任何特定任务头的模型架构

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

请注意，为 Hub 上的模型提供 --task 参数将禁用自动任务检测。

生成的 model.onnx 文件随后可以在支持 ONNX 标准的许多加速器之一上运行。例如，我们可以使用 optimum.onnxruntime 包加载和运行 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)

打印输出将给出以下结果

QuestionAnsweringModelOutput(loss=None, start_logits=tensor([[-4.7652, -1.0452, -7.0409, -4.6864, -4.0277, -6.2021, -4.9473,  2.6287,
          7.6111, -1.2488, -2.0551, -0.9350,  4.9758, -0.7707,  2.1493, -2.0703,
         -4.3232, -4.9472]]), end_logits=tensor([[ 0.4382, -1.6502, -6.3654, -6.0661, -4.1482, -3.5779, -0.0774, -3.6168,
         -1.8750, -2.8910,  6.2582,  0.5425, -3.7699,  3.8232, -1.5073,  6.2311,
          3.3604, -0.0772]]), hidden_states=None, attentions=None)

如您所见，将模型转换为 ONNX 并不意味着离开 Hugging Face 生态系统。您最终会得到与常规 🤗 Transformers 模型类似的 API！

也可以通过执行以下操作，直接从 ORTModelForQuestionAnswering 类将模型导出为 ONNX

>>> model = ORTModelForQuestionAnswering.from_pretrained("distilbert-base-uncased-distilled-squad", export=True)

有关更多信息，请查看 optimum.onnxruntime 文档关于此主题的页面。

对于 Hub 上的 TensorFlow 检查点，该过程是相同的。例如，我们可以从 Keras 组织导出纯 TensorFlow 检查点，如下所示

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

导出模型以用于 Optimum 的 ORTModel

通过 optimum-cli export onnx 导出的模型可以直接在ORTModel中使用。这对于编码器-解码器模型尤其有用，在这种情况下，导出会将编码器和解码器拆分为两个 .onnx 文件，因为编码器通常只运行一次，而解码器在自动生成任务中可能会运行多次。

在解码器中使用过去的键/值导出模型

当导出用于生成的解码器模型时，将过去键和值的重用封装在导出的 ONNX 中可能很有用。这允许避免在生成期间重新计算相同的中间激活。

在 ONNX 导出中，默认情况下会重用过去的键/值。此行为对应于 --task text2text-generation-with-past、--task text-generation-with-past 或 --task automatic-speech-recognition-with-past。如果出于任何目的您想禁用使用过去的键/值重用的导出，则需要显式传递给 optimum-cli export onnx 任务 text2text-generation、text-generation 或 automatic-speech-recognition。

使用过去的键/值导出的模型可以直接在 Optimum 的 ORTModel 中重用

optimum-cli export onnx --model gpt2 gpt2_onnx/

和

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

>>> tokenizer = AutoTokenizer.from_pretrained("./gpt2_onnx/")
>>> model = ORTModelForCausalLM.from_pretrained("./gpt2_onnx/")
>>> inputs = tokenizer("My name is Arthur and I live in", return_tensors="pt")
>>> gen_tokens = model.generate(**inputs)
>>> print(tokenizer.batch_decode(gen_tokens))
# prints ['My name is Arthur and I live in the United States of America. I am a member of the']

选择任务

从 Hugging Face Hub 上的模型导出时，在大多数情况下，不需要指定 --task。

但是，如果您需要检查给定模型架构 ONNX 导出支持哪些任务，我们已为您考虑周全。首先，您可以在此处查看 PyTorch 和 TensorFlow 支持的任务列表。

对于每个模型架构，您可以通过TasksManager找到支持的任务列表。例如，对于 DistilBERT，对于 ONNX 导出，我们有

>>> from optimum.exporters.tasks import TasksManager

>>> distilbert_tasks = list(TasksManager.get_supported_tasks_for_model_type("distilbert", "onnx").keys())
>>> print(distilbert_tasks)
['default', 'fill-mask', 'text-classification', 'multiple-choice', 'token-classification', 'question-answering']

然后，您可以将这些任务之一传递给 optimum-cli export onnx 命令中的 --task 参数，如上所述。

Transformers 模型的自定义导出

自定义官方 Transformers 模型的导出

Optimum 允许高级用户更精细地控制 ONNX 导出的配置。如果您想使用不同的关键字参数导出模型，例如使用 output_attentions=True 或 output_hidden_states=True，这将特别有用。

为了支持这些用例，~exporters.main_export 支持两个参数：model_kwargs 和 custom_onnx_configs，它们的使用方式如下

model_kwargs 允许覆盖模型 forward 的一些默认参数，实际上类似于 model(**reference_model_inputs, **model_kwargs)。
custom_onnx_configs 应该是一个 Dict[str, OnnxConfig]，从子模型名称（通常是 model、encoder_model、decoder_model 或 decoder_model_with_past - 参考）映射到给定子模型的自定义 ONNX 配置。

下面给出了一个完整的示例，允许导出 output_attentions=True 的模型。

from optimum.exporters.onnx import main_export
from optimum.exporters.onnx.model_configs import WhisperOnnxConfig
from transformers import AutoConfig

from optimum.exporters.onnx.base import ConfigBehavior
from typing import Dict

class CustomWhisperOnnxConfig(WhisperOnnxConfig):
    @property
    def outputs(self) -> Dict[str, Dict[int, str]]:
        common_outputs = super().outputs

        if self._behavior is ConfigBehavior.ENCODER:
            for i in range(self._config.encoder_layers):
                common_outputs[f"encoder_attentions.{i}"] = {0: "batch_size"}
        elif self._behavior is ConfigBehavior.DECODER:
            for i in range(self._config.decoder_layers):
                common_outputs[f"decoder_attentions.{i}"] = {
                    0: "batch_size",
                    2: "decoder_sequence_length",
                    3: "past_decoder_sequence_length + 1"
                }
            for i in range(self._config.decoder_layers):
                common_outputs[f"cross_attentions.{i}"] = {
                    0: "batch_size",
                    2: "decoder_sequence_length",
                    3: "encoder_sequence_length_out"
                }

        return common_outputs

    @property
    def torch_to_onnx_output_map(self):
        if self._behavior is ConfigBehavior.ENCODER:
            # The encoder export uses WhisperEncoder that returns the key "attentions"
            return {"attentions": "encoder_attentions"}
        else:
            return {}

model_id = "openai/whisper-tiny.en"
config = AutoConfig.from_pretrained(model_id)

custom_whisper_onnx_config = CustomWhisperOnnxConfig(
        config=config,
        task="automatic-speech-recognition",
)

encoder_config = custom_whisper_onnx_config.with_behavior("encoder")
decoder_config = custom_whisper_onnx_config.with_behavior("decoder", use_past=False)
decoder_with_past_config = custom_whisper_onnx_config.with_behavior("decoder", use_past=True)

custom_onnx_configs={
    "encoder_model": encoder_config,
    "decoder_model": decoder_config,
    "decoder_with_past_model": decoder_with_past_config,
}

main_export(
    model_id,
    output="custom_whisper_onnx",
    no_post_process=True,
    model_kwargs={"output_attentions": True},
    custom_onnx_configs=custom_onnx_configs
)

对于只需要单个 ONNX 文件的任务（例如，仅编码器），具有自定义输入/输出的导出模型随后可以与类 optimum.onnxruntime.ORTModelForCustomTasks 一起使用，以便在 CPU 或 GPU 上使用 ONNX Runtime 进行推理。

使用自定义建模自定义 Transformers 模型的导出

Optimum 支持导出使用 trust_remote_code=True 的自定义建模的 Transformers 模型，这些模型在 Transormers 库中未正式支持，但可以与其功能一起使用，如 pipelines 和 generation。

此类模型的示例包括 THUDM/chatglm2-6b 和 mosaicml/mpt-30b。

要导出自定义模型，需要将字典 custom_onnx_configs 传递给 main_export()，其中包含要导出的模型的所有子部分的 ONNX 配置定义（例如，编码器和解码器子部分）。以下示例允许导出 mosaicml/mpt-7b 模型

from optimum.exporters.onnx import main_export

from transformers import AutoConfig

from optimum.exporters.onnx.config import TextDecoderOnnxConfig
from optimum.utils import NormalizedTextConfig, DummyPastKeyValuesGenerator
from typing import Dict


class MPTDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
    """
    MPT swaps the two last dimensions for the key cache compared to usual transformers
    decoder models, thus the redefinition here.
    """
    def generate(self, input_name: str, framework: str = "pt"):
        past_key_shape = (
            self.batch_size,
            self.num_attention_heads,
            self.hidden_size // self.num_attention_heads,
            self.sequence_length,
        )
        past_value_shape = (
            self.batch_size,
            self.num_attention_heads,
            self.sequence_length,
            self.hidden_size // self.num_attention_heads,
        )
        return [
            (
                self.random_float_tensor(past_key_shape, framework=framework),
                self.random_float_tensor(past_value_shape, framework=framework),
            )
            for _ in range(self.num_layers)
        ]

class CustomMPTOnnxConfig(TextDecoderOnnxConfig):
    DUMMY_INPUT_GENERATOR_CLASSES = (MPTDummyPastKeyValuesGenerator,) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES
    DUMMY_PKV_GENERATOR_CLASS = MPTDummyPastKeyValuesGenerator

    DEFAULT_ONNX_OPSET = 14  # aten::tril operator requires opset>=14
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(
        hidden_size="d_model",
        num_layers="n_layers",
        num_attention_heads="n_heads"
    )

    def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str):
        """
        Adapted from https://github.com/huggingface/optimum/blob/v1.9.0/optimum/exporters/onnx/base.py#L625
        """
        if direction not in ["inputs", "outputs"]:
            raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')

        if direction == "inputs":
            decoder_sequence_name = "past_sequence_length"
            name = "past_key_values"
        else:
            decoder_sequence_name = "past_sequence_length + 1"
            name = "present"

        for i in range(self._normalized_config.num_layers):
            inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size", 3: decoder_sequence_name}
            inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size", 2: decoder_sequence_name}


model_id = "fxmarty/tiny-mpt-random-remote-code"
config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)

onnx_config = CustomMPTOnnxConfig(
    config=config,
    task="text-generation",
    use_past_in_inputs=False,
)
onnx_config_with_past = CustomMPTOnnxConfig(config, task="text-generation", use_past=True)

custom_onnx_configs = {
    "decoder_model": onnx_config,
    "decoder_with_past_model": onnx_config_with_past,
}

main_export(
    model_id,
    output="mpt_onnx",
    task="text-generation-with-past",
    trust_remote_code=True,
    custom_onnx_configs=custom_onnx_configs,
    no_post_process=True,
    legacy=True,
    opset=14
)

此外，main_export 的高级参数 fn_get_submodels 允许自定义在需要将模型导出到多个子模型时如何提取子模型。此类函数的示例可以在[此处查阅](link to utils.py relevant code once merged)。

< > 在 GitHub 上更新