#
Exportando modelos para ONNX

Se você precisar implantar modelos 🤗 Transformers em ambientes de produção, recomendamos exporta-los para um formato serializado que pode ser carregado e executado em tempos de execução e hardware. Neste guia, mostraremos como exportar modelos 🤗 Transformers para ONNX (Open Neural Network eXchange).

Uma vez exportado, um modelo pode ser otimizado para inferência por meio de técnicas como quantização e poda. Se você estiver interessado em otimizar seus modelos para serem executados com máxima eficiência, confira a biblioteca 🤗 Optimum .

ONNX é um padrão aberto que define um conjunto comum de operadores e um formato de arquivo comum para representar modelos de aprendizado profundo em uma ampla variedade de estruturas, incluindo PyTorch e TensorFlow. Quando um modelo é exportado para o formato ONNX, esses operadores são usados para construir um grafo computacional (muitas vezes chamado de representação intermediária) que representa o fluxo de dados através da rede neural.

Ao expor um grafo com operadores e tipos de dados padronizados, o ONNX facilita a alternar entre os frameworks. Por exemplo, um modelo treinado em PyTorch pode ser exportado para formato ONNX e depois importado no TensorFlow (e vice-versa).

🤗 Transformers fornece um pacote `transformers.onnx` que permite que você converta os checkpoints do modelo em um grafo ONNX aproveitando os objetos de configuração. Esses objetos de configuração vêm prontos para várias arquiteturas de modelo e são projetado para ser facilmente extensível a outras arquiteturas.

As configurações prontas incluem as seguintes arquiteturas:

- ALBERT
- BART
- BEiT
- BERT
- BigBird
- BigBird-Pegasus
- Blenderbot
- BlenderbotSmall
- BLOOM
- CamemBERT
- CLIP
- CodeGen
- Conditional DETR
- ConvBERT
- ConvNeXT
- ConvNeXTV2
- Data2VecText
- Data2VecVision
- DeBERTa
- DeBERTa-v2
- DeiT
- DETR
- DistilBERT
- ELECTRA
- ERNIE
- FlauBERT
- GPT Neo
- GPT-J
- GroupViT
- I-BERT
- LayoutLM
- LayoutLMv3
- LeViT
- Longformer
- LongT5
- M2M100
- Marian
- mBART
- MobileBERT
- MobileViT
- MT5
- OpenAI GPT-2
- OWL-ViT
- Perceiver
- PLBart
- ResNet
- RoBERTa
- RoFormer
- SegFormer
- SqueezeBERT
- Swin Transformer
- T5
- Table Transformer
- Vision Encoder decoder
- ViT
- XLM
- XLM-RoBERTa
- XLM-RoBERTa-XL
- YOLOS

Nas próximas duas seções, mostraremos como:

- Exportar um modelo suportado usando o pacote `transformers.onnx`.
- Exportar um modelo personalizado para uma arquitetura sem suporte.

##
Exportando um modelo para ONNX

Para exportar um modelo 🤗 Transformers para o ONNX, primeiro você precisa instalar algumas dependências extras:

    pip install transformers[onnx]

O pacote `transformers.onnx` pode então ser usado como um módulo Python:

    python -m transformers.onnx --help

    usage: Hugging Face Transformers ONNX exporter [-h] -m MODEL [--feature {causal-lm, ...}] [--opset OPSET] [--atol ATOL] output

    positional arguments:
      output                Path indicating where to store generated ONNX model.

    optional arguments:
      -h, --help            show this help message and exit
      -m MODEL, --model MODEL
                            Model ID on huggingface.co or path on disk to load model from.
      --feature {causal-lm, ...}
                            The type of features to export the model with.
      --opset OPSET         ONNX opset version to export the model with.
      --atol ATOL           Absolute difference tolerance when validating the model.

A exportação de um checkpoint usando uma configuração pronta pode ser feita da seguinte forma:

    python -m transformers.onnx --model=distilbert/distilbert-base-uncased onnx/

Você deve ver os seguintes logs:

    Validating ONNX model...
            -[✓] ONNX model output names match reference model ({'last_hidden_state'})
            - Validating ONNX Model output "last_hidden_state":
                    -[✓] (2, 8, 768) matches (2, 8, 768)
                    -[✓] all values close (atol: 1e-05)
    All good, model saved at: onnx/model.onnx

Isso exporta um grafo ONNX do ponto de verificação definido pelo argumento `--model`. Nisso Por exemplo, é `distilbert/distilbert-base-uncased`, mas pode ser qualquer checkpoint no Hugging Face Hub ou um armazenado localmente.

O arquivo `model.onnx` resultante pode ser executado em um dos muitos aceleradores que suportam o ONNX padrão. Por exemplo, podemos carregar e executar o modelo com ONNX Tempo de execução da seguinte forma:

    >>> 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))

Os nomes de saída necessários (como `["last_hidden_state"]`) podem ser obtidos pegando uma configuração ONNX de cada modelo. Por exemplo, para DistilBERT temos:

    >>> from transformers.models.distilbert import DistilBertConfig, DistilBertOnnxConfig

    >>> config = DistilBertConfig()
    >>> onnx_config = DistilBertOnnxConfig(config)
    >>> print(list(onnx_config.outputs.keys()))
    ["last_hidden_state"]

O processo é idêntico para os checkpoints do TensorFlow no Hub. Por exemplo, podemos exportar um checkpoint TensorFlow puro do Keras da seguinte forma:

    python -m transformers.onnx --model=keras-io/transformers-qa onnx/

Para exportar um modelo armazenado localmente, você precisará ter os pesos e arquivos tokenizer armazenados em um diretório. Por exemplo, podemos carregar e salvar um checkpoint como:

    >>> from transformers import AutoTokenizer, AutoModelForSequenceClassification

    >>> # Load tokenizer and PyTorch weights form the Hub
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
    >>> pt_model = AutoModelForSequenceClassification.from_pretrained("distilbert/distilbert-base-uncased")
    >>> # Save to disk
    >>> tokenizer.save_pretrained("local-pt-checkpoint")
    >>> pt_model.save_pretrained("local-pt-checkpoint")

Uma vez que o checkpoint é salvo, podemos exportá-lo para o ONNX apontando o `--model` argumento do pacote `transformers.onnx` para o diretório desejado:

    python -m transformers.onnx --model=local-pt-checkpoint onnx/

    >>> from transformers import AutoTokenizer, TFAutoModelForSequenceClassification

    >>> # Load tokenizer and TensorFlow weights from the Hub
    >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
    >>> tf_model = TFAutoModelForSequenceClassification.from_pretrained("distilbert/distilbert-base-uncased")
    >>> # Save to disk
    >>> tokenizer.save_pretrained("local-tf-checkpoint")
    >>> tf_model.save_pretrained("local-tf-checkpoint")

Uma vez que o checkpoint é salvo, podemos exportá-lo para o ONNX apontando o `--model` argumento do pacote `transformers.onnx` para o diretório desejado:

    python -m transformers.onnx --model=local-tf-checkpoint onnx/

##
Selecionando features para diferentes tarefas do modelo

Cada configuração pronta vem com um conjunto de features que permitem exportar modelos para diferentes tipos de tarefas. Conforme mostrado na tabela abaixo, cada recurso é associado a uma `AutoClass` diferente:

Feature | Auto Class
--- | ---
`causal-lm`, `causal-lm-with-past` | `AutoModelForCausalLM`
`default`, `default-with-past` | `AutoModel`
`masked-lm` | `AutoModelForMaskedLM`
`question-answering` | `AutoModelForQuestionAnswering`
`seq2seq-lm`, `seq2seq-lm-with-past` | `AutoModelForSeq2SeqLM`
`sequence-classification` | `AutoModelForSequenceClassification`
`token-classification` | `AutoModelForTokenClassification`

Para cada configuração, você pode encontrar a lista de recursos suportados por meio do `FeaturesManager`. Por exemplo, para DistilBERT temos:

    >>> from transformers.onnx.features import FeaturesManager

    >>> distilbert_features = list(FeaturesManager.get_supported_features_for_model_type("distilbert").keys())
    >>> print(distilbert_features)
    ["default", "masked-lm", "causal-lm", "sequence-classification", "token-classification", "question-answering"]

Você pode então passar um desses recursos para o argumento `--feature` no pacote `transformers.onnx`. Por exemplo, para exportar um modelo de classificação de texto, podemos escolher um modelo ajustado no Hub e executar:

    python -m transformers.onnx --model=distilbert/distilbert-base-uncased-finetuned-sst-2-english \
                                --feature=sequence-classification onnx/

Isso exibe os seguintes logs:

    Validating ONNX model...
            -[✓] ONNX model output names match reference model ({'logits'})
            - Validating ONNX Model output "logits":
                    -[✓] (2, 2) matches (2, 2)
                    -[✓] all values close (atol: 1e-05)
    All good, model saved at: onnx/model.onnx

Observe que, neste caso, os nomes de saída do modelo ajustado são `logits` em vez do `last_hidden_state` que vimos com o checkpoint `distilbert/distilbert-base-uncased` mais cedo. Isso é esperado, pois o modelo ajustado (fine-tuned) possui uma cabeça de classificação de sequência.

Os recursos que têm um sufixo `with-pass` (como `causal-lm-with-pass`) correspondem a classes de modelo com estados ocultos pré-computados (chave e valores nos blocos de atenção) que pode ser usado para decodificação autorregressiva rápida.

Para modelos do tipo `VisionEncoderDecoder`, as partes do codificador e do decodificador são exportados separadamente como dois arquivos ONNX chamados `encoder_model.onnx` e `decoder_model.onnx` respectivamente.

##
Exportando um modelo para uma arquitetura sem suporte

Se você deseja exportar um modelo cuja arquitetura não é suportada nativamente pela biblioteca, há três etapas principais a seguir:

1. Implemente uma configuração ONNX personalizada.
2. Exporte o modelo para o ONNX.
3. Valide as saídas do PyTorch e dos modelos exportados.

Nesta seção, veremos como o DistilBERT foi implementado para mostrar o que está envolvido em cada passo.

###
Implementando uma configuração ONNX personalizada

Vamos começar com o objeto de configuração ONNX. Fornecemos três classes abstratas que você deve herdar, dependendo do tipo de arquitetura de modelo que deseja exportar:

- Modelos baseados em codificador herdam de `OnnxConfig`
- Modelos baseados em decodificador herdam de `OnnxConfigWithPast`
- Os modelos codificador-decodificador herdam de `OnnxSeq2SeqConfigWithPast`

Uma boa maneira de implementar uma configuração ONNX personalizada é observar as implementação no arquivo `configuration_<model_name>.py` de uma arquitetura semelhante.

Como o DistilBERT é um modelo baseado em codificador, sua configuração é herdada de `OnnxConfig`:

    >>> from typing import Mapping, OrderedDict
    >>> from transformers.onnx import OnnxConfig

    >>> class DistilBertOnnxConfig(OnnxConfig):
    ...     @property
    ...     def inputs(self) -> Mapping[str, Mapping[int, str]]:
    ...         return OrderedDict(
    ...             [
    ...                 ("input_ids", {0: "batch", 1: "sequence"}),
    ...                 ("attention_mask", {0: "batch", 1: "sequence"}),
    ...             ]
    ...         )

Todo objeto de configuração deve implementar a propriedade `inputs` e retornar um mapeamento, onde cada chave corresponde a uma entrada esperada e cada valor indica o eixo dessa entrada. Para o DistilBERT, podemos ver que duas entradas são necessárias: `input_ids` e `attention_mask`. Essas entradas têm a mesma forma de `(batch_size, sequence_length)` é por isso que vemos os mesmos eixos usados na configuração.

Notice that `inputs` property for `DistilBertOnnxConfig` returns an `OrderedDict`. This ensures that the inputs are matched with their relative position within the `PreTrainedModel.forward()` method when tracing the graph. We recommend using an `OrderedDict` for the `inputs` and `outputs` properties when implementing custom ONNX configurations.

Observe que a propriedade `inputs` para `DistilBertOnnxConfig` retorna um `OrderedDict`. Este garante que as entradas sejam combinadas com sua posição relativa dentro do método `PreTrainedModel.forward()` ao traçar o grafo. Recomendamos o uso de um `OrderedDict` para as propriedades `inputs` e `outputs` ao implementar configurações personalizadas ONNX.

Depois de implementar uma configuração ONNX, você pode instanciá-la fornecendo a configuração do modelo base da seguinte forma:

    >>> from transformers import AutoConfig

    >>> config = AutoConfig.from_pretrained("distilbert/distilbert-base-uncased")
    >>> onnx_config = DistilBertOnnxConfig(config)

O objeto resultante tem várias propriedades úteis. Por exemplo, você pode visualizar o conjunto de operadores ONNX que será usado durante a exportação:

    >>> print(onnx_config.default_onnx_opset)
    11

Você também pode visualizar as saídas associadas ao modelo da seguinte forma:

    >>> print(onnx_config.outputs)
    OrderedDict([("last_hidden_state", {0: "batch", 1: "sequence"})])

Observe que a propriedade outputs segue a mesma estrutura das entradas; ele retorna um `OrderedDict` de saídas nomeadas e suas formas. A estrutura de saída está ligada a escolha do recurso com o qual a configuração é inicializada. Por padrão, a configuração do ONNX é inicializada com o recurso `default` que corresponde à exportação de um modelo carregado com a classe `AutoModel`. Se você deseja exportar um modelo para outra tarefa, apenas forneça um recurso diferente para o argumento `task` quando você inicializar a configuração ONNX . Por exemplo, se quisermos exportar o DistilBERT com uma sequência de classificação, poderíamos usar:

    >>> from transformers import AutoConfig

    >>> config = AutoConfig.from_pretrained("distilbert/distilbert-base-uncased")
    >>> onnx_config_for_seq_clf = DistilBertOnnxConfig(config, task="sequence-classification")
    >>> print(onnx_config_for_seq_clf.outputs)
    OrderedDict([('logits', {0: 'batch'})])

Todas as propriedades e métodos básicos associados a `OnnxConfig` e as outras classes de configuração podem ser substituídas se necessário. Confira `BartOnnxConfig` para um exemplo avançado.

###
Exportando um modelo

Depois de ter implementado a configuração do ONNX, o próximo passo é exportar o modelo. Aqui podemos usar a função `export()` fornecida pelo pacote `transformers.onnx`. Esta função espera a configuração do ONNX, juntamente com o modelo base e o tokenizer, e o caminho para salvar o arquivo exportado:

    >>> from pathlib import Path
    >>> from transformers.onnx import export
    >>> from transformers import AutoTokenizer, AutoModel

    >>> onnx_path = Path("model.onnx")
    >>> model_ckpt = "distilbert/distilbert-base-uncased"
    >>> base_model = AutoModel.from_pretrained(model_ckpt)
    >>> tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

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

Os `onnx_inputs` e `onnx_outputs` retornados pela função `export()` são listas de chaves definidas nas propriedades `inputs` e `outputs` da configuração. Uma vez que o modelo é exportado, você pode testar se o modelo está bem formado da seguinte forma:

    >>> import onnx

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

Se o seu modelo for maior que 2GB, você verá que muitos arquivos adicionais são criados durante a exportação. Isso é esperado porque o ONNX usa Protocol Buffers para armazenar o modelo e estes têm um limite de tamanho de 2GB. Veja a ONNX documentação para instruções sobre como carregar modelos com dados externos.

###
Validando a saída dos modelos

A etapa final é validar se as saídas do modelo base e exportado concordam dentro de alguma tolerância absoluta. Aqui podemos usar a função `validate_model_outputs()` fornecida pelo pacote `transformers.onnx` da seguinte forma:

    >>> from transformers.onnx import validate_model_outputs

    >>> validate_model_outputs(
    ...     onnx_config, tokenizer, base_model, onnx_path, onnx_outputs, onnx_config.atol_for_validation
    ... )

Esta função usa o método `generate_dummy_inputs()` para gerar entradas para o modelo base e o exportado, e a tolerância absoluta pode ser definida na configuração. Geralmente encontramos concordância numérica em 1e-6 a 1e-4 de alcance, embora qualquer coisa menor que 1e-3 provavelmente esteja OK.

##
Contribuindo com uma nova configuração para 🤗 Transformers

Estamos procurando expandir o conjunto de configurações prontas e receber contribuições da comunidade! Se você gostaria de contribuir para a biblioteca, você precisará:

- Implemente a configuração do ONNX no arquivo `configuration_<model_name>.py` correspondente Arquivo
- Incluir a arquitetura do modelo e recursos correspondentes em `~onnx.features.FeatureManager`
- Adicione sua arquitetura de modelo aos testes em `test_onnx_v2.py`

Confira como ficou a configuração do IBERT para obter uma idéia do que está envolvido.

< > Update on GitHub