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
[~transformers.onnx.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:
- Implemente uma configuração ONNX personalizada.
- Exporte o modelo para o ONNX.
- 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 [
~onnx.config.OnnxConfig
] - Modelos baseados em decodificador herdam de [
~onnx.config.OnnxConfigWithPast
] - Os modelos codificador-decodificador herdam de [
~onnx.config.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 [~onnx.config.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 [~transformers.onnx.OnnxConfig.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.