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

摘要

将模型导出到 ONNX 就像

optimum-cli export onnx --model gpt2 gpt2_onnx/

查看帮助以获取更多选项

optimum-cli export onnx --help

为什么要使用 ONNX？

如果你需要在生产环境中部署 🤗 Transformers 或 🤗 Diffusers 模型，我们建议将它们导出为可加载并在专用运行时和硬件上执行的序列化格式。在本指南中，我们将向你展示如何将这些模型导出到 ONNX (Open Neural Network eXchange)。

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。如果出于任何目的你希望禁用带有历史键/值重用的导出，则需要明确将任务 text2text-generation、text-generation 或 automatic-speech-recognition 传递给 optimum-cli export onnx。

使用历史键/值导出的模型可以直接在 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 模型，这在 Transformers 库中不被官方支持，但可与其功能一起使用，如管道和生成。

此类模型的示例包括 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 允许在模型需要导出为多个子模型时自定义子模型的提取方式。此类函数的示例可[在此处查阅](合并后链接到 utils.py 相关代码)。

< > 在 GitHub 上更新