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

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:

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.

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

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.

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'})])

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)

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.