添加对不受支持架构的支持

如果您希望导出库尚不支持其架构的模型，则需要遵循以下主要步骤

实现自定义 ONNX 配置。
在任务管理器中注册 ONNX 配置。
将模型导出到 ONNX。
验证原始模型和导出模型的输出。

在本节中，我们将研究 BERT 的实现方式，以展示每个步骤涉及的内容。

实现自定义 ONNX 配置

让我们从 ONNX 配置对象开始。我们提供了一个三级类层次结构，并且为了添加对模型的支持，大多数情况下，从正确的中间端类继承将是最佳选择。如果您正在添加处理以前从未见过的模态和/或情况的架构，则可能需要自己实现中间端类。

实现自定义 ONNX 配置的一个好方法是查看 optimum/exporters/onnx/model_configs.py 文件中现有的配置实现。

此外，如果您尝试添加的架构与已支持的架构（例如，在已支持 BERT 的情况下添加对 ALBERT 的支持）非常相似，则尝试简单地从该类继承可能会起作用。

从中间端类继承时，请查找处理与您尝试支持的模型相同的模态/类别的类。

示例：添加对 BERT 的支持

由于 BERT 是一个基于编码器的文本模型，因此其配置继承自中间端类 TextEncoderOnnxConfig。在 optimum/exporters/onnx/model_configs.py 中

# This class is actually in optimum/exporters/onnx/config.py
class TextEncoderOnnxConfig(OnnxConfig):
    # Describes how to generate the dummy inputs.
    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,)

class BertOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4

    @property
    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }

首先，让我们解释一下 TextEncoderOnnxConfig 的作用。虽然大多数功能已经在 OnnxConfig 中实现，但此类与模态无关，这意味着它不知道应该处理哪种输入。输入生成的方式是通过 DUMMY_INPUT_GENERATOR_CLASSES 属性，它是一个 DummyInputGenerator 的元组。在这里，我们通过指定 DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) 来创建一个从 OnnxConfig 继承的模态感知配置。

然后是模型特定的类 BertOnnxConfig。这里指定了两个类属性

NORMALIZED_CONFIG_CLASS：这必须是一个 NormalizedConfig，它基本上允许输入生成器以通用方式访问模型配置属性。
ATOL_FOR_VALIDATION：它在使用原始模型验证导出模型时使用，这是输出值差异的绝对可接受容差。

每个配置对象都必须实现 inputs 属性并返回一个映射，其中每个键对应于一个输入名称，每个值指示该输入中哪些轴是动态的。对于 BERT，我们可以看到需要三个输入：input_ids、attention_mask 和 token_type_ids。这些输入具有相同的 (batch_size, sequence_length) 形状（除了 多项选择 任务），这就是我们在配置中看到使用相同轴的原因。

实现 ONNX 配置后，您可以通过提供基础模型的配置来实例化它，如下所示

>>> from transformers import AutoConfig
>>> from optimum.exporters.onnx.model_configs import BertOnnxConfig
>>> config = AutoConfig.from_pretrained("bert-base-uncased")
>>> onnx_config = BertOnnxConfig(config)

生成的具有几个有用的属性。例如，您可以查看导出过程中将使用的 ONNX 运算符集

>>> print(onnx_config.DEFAULT_ONNX_OPSET)
11

您还可以查看与模型关联的输出，如下所示

>>> print(onnx_config.outputs)
OrderedDict([('last_hidden_state', {0: 'batch_size', 1: 'sequence_length'})])

请注意，outputs 属性与 inputs 具有相同的结构；它返回一个命名输出及其形状的 OrderedDict。输出结构与配置初始化时选择的任务相关联。默认情况下，ONNX 配置使用 default 任务初始化，该任务对应于使用 AutoModel 类加载的模型的导出。如果您想为其他任务导出模型，只需在初始化 ONNX 配置时向 task 参数提供不同的任务即可。例如，如果我们希望使用序列分类头导出 BERT，我们可以使用

>>> from transformers import AutoConfig

>>> config = AutoConfig.from_pretrained("bert-base-uncased")
>>> onnx_config_for_seq_clf = BertOnnxConfig(config, task="text-classification")
>>> print(onnx_config_for_seq_clf.outputs)
OrderedDict([('logits', {0: 'batch_size'})])

查看 BartOnnxConfig 以获取高级示例。

在任务管理器中注册 ONNX 配置

任务管理器是给定名称和任务加载模型并获取给定（架构、后端）组合的正确配置的主要入口点。在添加对导出到 ONNX 的支持时，将配置注册到 TasksManager 将使导出在命令行工具中可用。

为此，请在_SUPPORTED_MODEL_TYPE属性中添加一个条目。

如果该模型已在 ONNX 以外的其他后端得到支持，则它将已经有一个条目，因此您只需要添加一个onnx键，并指定配置类的名称。
否则，您需要添加整个条目。

对于 BERT，它看起来如下所示：

    "bert": supported_tasks_mapping(
        "default",
        "fill-mask",
        "text-generation",
        "text-classification",
        "multiple-choice",
        "token-classification",
        "question-answering",
        onnx="BertOnnxConfig",
    )

导出模型

一旦您实现了 ONNX 配置，下一步就是导出模型。在这里，我们可以使用optimum.exporters.onnx包提供的export()函数。此函数需要 ONNX 配置、基础模型以及保存导出文件的路径。

>>> from pathlib import Path
>>> from optimum.exporters import TasksManager
>>> from optimum.exporters.onnx import export
>>> from transformers import AutoModel

>>> base_model = AutoModel.from_pretrained("bert-base-uncased")

>>> onnx_path = Path("model.onnx")
>>> onnx_config_constructor = TasksManager.get_exporter_config_constructor("onnx", base_model)
>>> onnx_config = onnx_config_constructor(base_model.config)

>>> onnx_inputs, onnx_outputs = export(base_model, onnx_config, onnx_path, onnx_config.DEFAULT_ONNX_OPSET)

export()函数返回的onnx_inputs和onnx_outputs是配置的inputs和inputs属性中定义的键的列表。模型导出后，您可以按照以下步骤测试模型是否格式正确。

>>> import onnx

>>> onnx_model = onnx.load("model.onnx")
>>> onnx.checker.check_model(onnx_model)

如果您的模型大于 2GB，您会发现导出过程中会创建许多其他文件。这是预期的，因为 ONNX 使用Protocol Buffers来存储模型，并且这些文件的大小限制为 2GB。有关如何加载具有外部数据的模型的说明，请参阅ONNX 文档。

验证模型输出

最后一步是验证基础模型和导出模型的输出在某个绝对容差范围内是否一致。在这里，我们可以使用optimum.exporters.onnx包提供的validate_model_outputs()函数。

>>> from optimum.exporters.onnx import validate_model_outputs

>>> validate_model_outputs(
...     onnx_config, base_model, onnx_path, onnx_outputs, onnx_config.ATOL_FOR_VALIDATION
... )

将新的配置贡献给 🤗 Optimum

现在，架构的支持已实现并得到验证，还剩下两件事。

将您的模型架构添加到tests/exporters/test_onnx_export.py中的测试中。
在optimum存储库上创建一个 PR。

感谢您的贡献！

< > 在 GitHub 上更新