新增對不受支援架構的支援

如果您希望匯出庫中尚未支援其架構的模型，請遵循以下主要步驟：

實現自定義 ONNX 配置。
在 TasksManager 中註冊 ONNX 配置。
將模型匯出到 ONNX。
驗證原始模型和匯出模型的輸出。

在本節中，我們將介紹 BERT 是如何實現的，以展示每個步驟所涉及的內容。

實現自定義 ONNX 配置

讓我們從 ONNX 配置物件開始。我們提供了一個 3 級類層次結構，為了新增對模型的支援，大多數時候都將從正確的中端類繼承。如果您要新增一個處理以前從未見過的模態和/或情況的架構，您可能必須自己實現一箇中端類。

實現自定義 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` 相同的結構；它返回一個有序字典，其中包含命名輸出及其形狀。輸出結構與配置初始化時選擇的任務相關聯。預設情況下，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` 以獲取高階示例。

在 TasksManager 中註冊 ONNX 配置

TasksManager 是根據名稱和任務載入模型，並獲取給定（架構，後端）對的正確配置的主要入口點。當新增對匯出到 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 上更新