Exportar modelos 🤗 Transformers

Si necesitas implementar modelos 🤗 Transformers en entornos de producción, te recomendamos exportarlos a un formato serializado que se pueda cargar y ejecutar en tiempos de ejecución y hardware especializados. En esta guía, te mostraremos cómo exportar modelos 🤗 Transformers en dos formatos ampliamente utilizados: ONNX y TorchScript.

Una vez exportado, un modelo puede optimizarse para la inferencia a través de técnicas como la cuantización y pruning. Si estás interesado en optimizar tus modelos para que funcionen con la máxima eficiencia, consulta la biblioteca de 🤗 Optimum.

ONNX

El proyecto ONNX (Open Neural Network eXchange) es un estándar abierto que define un conjunto común de operadores y un formato de archivo común para representar modelos de aprendizaje profundo en una amplia variedad de frameworks, incluidos PyTorch y TensorFlow. Cuando un modelo se exporta al formato ONNX, estos operadores se usan para construir un grafo computacional (a menudo llamado representación intermedia) que representa el flujo de datos a través de la red neuronal.

Al exponer un grafo con operadores y tipos de datos estandarizados, ONNX facilita el cambio entre frameworks. Por ejemplo, un modelo entrenado en PyTorch se puede exportar a formato ONNX y luego importar en TensorFlow (y viceversa).

🤗 Transformers proporciona un paquete llamado transformers.onnx, el cual permite convertir los checkpoints de un modelo en un grafo ONNX aprovechando los objetos de configuración. Estos objetos de configuración están hechos a la medida de diferentes arquitecturas de modelos y están diseñados para ser fácilmente extensibles a otras arquitecturas.

Las configuraciones a la medida incluyen las siguientes arquitecturas:

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

En las próximas dos secciones, te mostraremos cómo:

Exportar un modelo compatible utilizando el paquete transformers.onnx.
Exportar un modelo personalizado para una arquitectura no compatible.

Exportar un model a ONNX

Para exportar un modelo 🤗 Transformers a ONNX, tienes que instalar primero algunas dependencias extra:

pip install transformers[onnx]

El paquete transformers.onnx puede ser usado luego como un módulo de 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 tolerence when validating the model.

Exportar un checkpoint usando una configuración a la medida se puede hacer de la siguiente manera:

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

que debería mostrar los siguientes registros:

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

Esto exporta un grafo ONNX del checkpoint definido por el argumento --model. En este ejemplo, es un modelo distilbert/distilbert-base-uncased, pero puede ser cualquier checkpoint en Hugging Face Hub o que esté almacenado localmente.

El archivo model.onnx resultante se puede ejecutar en uno de los muchos aceleradores que admiten el estándar ONNX. Por ejemplo, podemos cargar y ejecutar el modelo con ONNX Runtime de la siguiente manera:

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

Los nombres necesarios de salida (es decir, ["last_hidden_state"]) se pueden obtener echando un vistazo a la configuración ONNX de cada modelo. Por ejemplo, para DistilBERT tenemos:

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

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

El proceso es idéntico para los checkpoints de TensorFlow en Hub. Por ejemplo, podemos exportar un checkpoint puro de TensorFlow desde Keras de la siguiente manera:

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

Para exportar un modelo que está almacenado localmente, deberás tener los pesos y tokenizadores del modelo almacenados en un directorio. Por ejemplo, podemos cargar y guardar un checkpoint de la siguiente manera:

```python >>> 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")


Una vez que se guarda el checkpoint, podemos exportarlo a ONNX usando el argumento `--model` 
del paquete `transformers.onnx` al directorio deseado:

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

```python >>> 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")


Una vez que se guarda el checkpoint, podemos exportarlo a ONNX usando el argumento `--model` 
del paquete `transformers.onnx` al directorio deseado:

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

Seleccionar características para diferentes topologías de un modelo

Cada configuración a la medida viene con un conjunto de características que te permiten exportar modelos para diferentes tipos de topologías o tareas. Como se muestra en la siguiente tabla, cada función está asociada con una auto-clase de automóvil 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 configuración, puedes encontrar la lista de funciones admitidas a través de FeaturesManager. Por ejemplo, para DistilBERT tenemos:

>>> 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"]

Le puedes pasar una de estas características al argumento --feature en el paquete transformers.onnx. Por ejemplo, para exportar un modelo de clasificación de texto, podemos elegir un modelo ya ajustado del Hub y ejecutar:

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

que mostrará los siguientes registros:

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

Ten en cuenta que, en este caso, los nombres de salida del modelo ajustado son logits en lugar de last_hidden_state que vimos anteriormente con el checkpoint distilbert/distilbert-base-uncased. Esto es de esperarse ya que el modelo ajustado tiene un cabezal de clasificación secuencial.

Las características que tienen un sufijo 'with-past' (por ejemplo, 'causal-lm-with-past') corresponden a topologías de modelo con estados ocultos precalculados (clave y valores en los bloques de atención) que se pueden usar para una decodificación autorregresiva más rápida.

Exportar un modelo para una arquitectura no compatible

Si deseas exportar un modelo cuya arquitectura no es compatible de forma nativa con la biblioteca, debes seguir tres pasos principales:

Implementa una configuración personalizada en ONNX.
Exporta el modelo a ONNX.
Valide los resultados de PyTorch y los modelos exportados.

En esta sección, veremos cómo se implementó la serialización de DistilBERT para mostrar lo que implica cada paso.

Implementar una configuración personalizada en ONNX

Comencemos con el objeto de configuración de ONNX. Proporcionamos tres clases abstractas de las que debe heredar, según el tipo de arquitectura del modelo que quieras exportar:

Modelos basados en el Encoder inherente de [~onnx.config.OnnxConfig]
Modelos basados en el Decoder inherente de [~onnx.config.OnnxConfigWithPast]
Modelos Encoder-decoder inherente de [~onnx.config.OnnxSeq2SeqConfigWithPast]

Una buena manera de implementar una configuración personalizada en ONNX es observar la implementación existente en el archivo configuration_<model_name>.py de una arquitectura similar.

Dado que DistilBERT es un modelo de tipo encoder, su configuración se hereda 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"}),
...             ]
...         )

Cada objeto de configuración debe implementar la propiedad inputs y devolver un mapeo, donde cada llave corresponde a una entrada esperada y cada valor indica el eje de esa entrada. Para DistilBERT, podemos ver que se requieren dos entradas: input_ids y attention_mask. Estas entradas tienen la misma forma de (batch_size, sequence_length), es por lo que vemos los mismos ejes utilizados en la configuración.

Observa que la propiedad inputs para DistilBertOnnxConfig devuelve un OrderedDict. Esto nos asegura que las entradas coincidan con su posición relativa dentro del método PreTrainedModel.forward() al rastrear el grafo. Recomendamos usar un OrderedDict para las propiedades inputs y outputs al implementar configuraciones ONNX personalizadas.

Una vez que hayas implementado una configuración ONNX, puedes crear una instancia proporcionando la configuración del modelo base de la siguiente manera:

>>> from transformers import AutoConfig

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

El objeto resultante tiene varias propiedades útiles. Por ejemplo, puedes ver el conjunto de operadores ONNX que se utilizará durante la exportación:

>>> print(onnx_config.default_onnx_opset)
11

También puedes ver los resultados asociados con el modelo de la siguiente manera:

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

Observa que la propiedad de salidas sigue la misma estructura que las entradas; devuelve un objecto OrderedDict de salidas nombradas y sus formas. La estructura de salida está vinculada a la elección de la función con la que se inicializa la configuración. Por defecto, la configuración de ONNX se inicializa con la función default que corresponde a exportar un modelo cargado con la clase AutoModel. Si quieres exportar una topología de modelo diferente, simplemente proporciona una característica diferente al argumento task cuando inicialices la configuración de ONNX. Por ejemplo, si quisiéramos exportar DistilBERT con un cabezal de clasificación de secuencias, podrí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 las propiedades base y métodos asociados con [~onnx.config.OnnxConfig] y las otras clases de configuración se pueden sobreescribir si es necesario. Consulte [BartOnnxConfig] para ver un ejemplo avanzado.

Exportar el modelo

Una vez que hayas implementado la configuración de ONNX, el siguiente paso es exportar el modelo. Aquí podemos usar la función export() proporcionada por el paquete transformers.onnx. Esta función espera la configuración de ONNX, junto con el modelo base y el tokenizador, y la ruta para guardar el archivo 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)

Los objetos onnx_inputs y onnx_outputs devueltos por la función export() son listas de llaves definidas en las propiedades inputs y outputs de la configuración. Una vez exportado el modelo, puedes probar que el modelo está bien formado de la siguiente manera:

>>> import onnx

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

Si tu modelo tiene más de 2GB, verás que se crean muchos archivos adicionales durante la exportación. Esto es esperado porque ONNX usa Búferes de protocolo para almacenar el modelo y éstos tienen un límite de tamaño de 2 GB. Consulta la documentación de ONNX para obtener instrucciones sobre cómo cargar modelos con datos externos.

Validar los resultados del modelo

El paso final es validar que los resultados del modelo base y exportado coincidan dentro de cierta tolerancia absoluta. Aquí podemos usar la función validate_model_outputs() proporcionada por el paquete transformers.onnx de la siguiente manera:

>>> 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 función usa el método OnnxConfig.generate_dummy_inputs() para generar entradas para el modelo base y exportado, y la tolerancia absoluta se puede definir en la configuración. En general, encontramos una concordancia numérica en el rango de 1e-6 a 1e-4, aunque es probable que cualquier valor menor que 1e-3 esté bien.

Contribuir con una nueva configuración a 🤗 Transformers

¡Estamos buscando expandir el conjunto de configuraciones a la medida para usar y agradecemos las contribuciones de la comunidad! Si deseas contribuir con su colaboración a la biblioteca, deberás:

Implementa la configuración de ONNX en el archivo configuration_<model_name>.py correspondiente
Incluye la arquitectura del modelo y las características correspondientes en [~onnx.features.FeatureManager]
Agrega tu arquitectura de modelo a las pruebas en test_onnx_v2.py

Revisa cómo fue la contribución para la configuración de IBERT y así tener una idea de lo que necesito.

TorchScript

Este es el comienzo de nuestros experimentos con TorchScript y todavía estamos explorando sus capacidades con modelos de tamaño de entrada variable. Es un tema de interés y profundizaremos nuestro análisis en las próximas versiones, con más ejemplos de código, una implementación más flexible y puntos de referencia que comparen códigos basados en Python con TorchScript compilado.

Según la documentación de PyTorch: "TorchScript es una forma de crear modelos serializables y optimizables a partir del código de PyTorch". Los dos módulos de Pytorch JIT y TRACE permiten al desarrollador exportar su modelo para reutilizarlo en otros programas, como los programas C++ orientados a la eficiencia.

Hemos proporcionado una interfaz que permite exportar modelos de 🤗 Transformers a TorchScript para que puedan reutilizarse en un entorno diferente al de un programa Python basado en PyTorch. Aquí explicamos cómo exportar y usar nuestros modelos usando TorchScript.

Exportar un modelo requiere de dos cosas:

un pase hacia adelante con entradas ficticias.
instanciación del modelo con la indicador torchscript.

Estas necesidades implican varias cosas con las que los desarrolladores deben tener cuidado. Éstas se detallan a continuación.

Indicador de TorchScript y pesos atados

Este indicador es necesario porque la mayoría de los modelos de lenguaje en este repositorio tienen pesos vinculados entre su capa de Embedding y su capa de Decoding. TorchScript no permite la exportación de modelos que tengan pesos atados, por lo que es necesario desvincular y clonar los pesos previamente.

Esto implica que los modelos instanciados con el indicador torchscript tienen su capa Embedding y Decoding separadas, lo que significa que no deben entrenarse más adelante. El entrenamiento desincronizaría las dos capas, lo que generaría resultados inesperados.

Este no es el caso de los modelos que no tienen un cabezal de modelo de lenguaje, ya que no tienen pesos atados. Estos modelos se pueden exportar de forma segura sin el indicador torchscript.

Entradas ficticias y longitudes estándar

Las entradas ficticias se utilizan para crear un modelo de pase hacia adelante. Mientras los valores de las entradas se propagan a través de las capas, PyTorch realiza un seguimiento de las diferentes operaciones ejecutadas en cada tensor. Estas operaciones registradas se utilizan luego para crear el "rastro" del modelo.

El rastro se crea en relación con las dimensiones de las entradas. Por lo tanto, está limitado por las dimensiones de la entrada ficticia y no funcionará para ninguna otra longitud de secuencia o tamaño de lote. Al intentar con un tamaño diferente, un error como:

The expanded size of the tensor (3) must match the existing size (7) at non-singleton dimension 2

aparecerá. Por lo tanto, se recomienda rastrear el modelo con un tamaño de entrada ficticia al menos tan grande como la entrada más grande que se alimentará al modelo durante la inferencia. El padding se puede realizar para completar los valores que faltan. Sin embargo, como el modelo se habrá rastreado con un tamaño de entrada grande, las dimensiones de las diferentes matrices también serán grandes, lo que dará como resultado más cálculos.

Se recomienda tener cuidado con el número total de operaciones realizadas en cada entrada y seguir de cerca el rendimiento al exportar modelos de longitud de secuencia variable.

Usar TorchScript en Python

A continuación se muestra un ejemplo que muestra cómo guardar, cargar modelos y cómo usar el rastreo para la inferencia.

Guardando un modelo

Este fragmento muestra cómo usar TorchScript para exportar un BertModel. Aquí, el BertModel se instancia de acuerdo con la clase BertConfig y luego se guarda en el disco con el nombre de archivo traced_bert.pt

from transformers import BertModel, BertTokenizer, BertConfig
import torch

enc = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")

# Tokenizing input text
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = enc.tokenize(text)

# Masking one of the input tokens
masked_index = 8
tokenized_text[masked_index] = "[MASK]"
indexed_tokens = enc.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]

# Creating a dummy input
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
dummy_input = [tokens_tensor, segments_tensors]

# Initializing the model with the torchscript flag
# Flag set to True even though it is not necessary as this model does not have an LM Head.
config = BertConfig(
    vocab_size_or_config_json_file=32000,
    hidden_size=768,
    num_hidden_layers=12,
    num_attention_heads=12,
    intermediate_size=3072,
    torchscript=True,
)

# Instantiating the model
model = BertModel(config)

# The model needs to be in evaluation mode
model.eval()

# If you are instantiating the model with *from_pretrained* you can also easily set the TorchScript flag
model = BertModel.from_pretrained("google-bert/bert-base-uncased", torchscript=True)

# Creating the trace
traced_model = torch.jit.trace(model, [tokens_tensor, segments_tensors])
torch.jit.save(traced_model, "traced_bert.pt")

Cargar un modelo

Este fragmento muestra cómo cargar el BertModel que se guardó previamente en el disco con el nombre traced_bert.pt. Estamos reutilizando el dummy_input previamente inicializado.

loaded_model = torch.jit.load("traced_bert.pt")
loaded_model.eval()

all_encoder_layers, pooled_output = loaded_model(*dummy_input)

Usar un modelo rastreado para la inferencia

Usar el modelo rastreado para la inferencia es tan simple como usar su método __call__:

traced_model(tokens_tensor, segments_tensors)

Implementar los modelos HuggingFace TorchScript en AWS mediante Neuron SDK

AWS presentó la familia de instancias Amazon EC2 Inf1 para la inferencia de aprendizaje automático de bajo costo y alto rendimiento en la nube. Las instancias Inf1 funcionan con el chip AWS Inferentia, un acelerador de hardware personalizado, que se especializa en cargas de trabajo de inferencia de aprendizaje profundo. AWS Neuron es el kit de desarrollo para Inferentia que admite el rastreo y la optimización de modelos de transformers para su implementación en Inf1. El SDK de Neuron proporciona:

API fácil de usar con una línea de cambio de código para rastrear y optimizar un modelo de TorchScript para la inferencia en la nube.
Optimizaciones de rendimiento listas para usar con un costo-rendimiento mejorado
Soporte para modelos HuggingFace Transformers construidos con PyTorch o TensorFlow.

Implicaciones

Los modelos Transformers basados en la arquitectura BERT (Representaciones de Enconder bidireccional de Transformers), o sus variantes, como distilBERT y roBERTa, se ejecutarán mejor en Inf1 para tareas no generativas, como la respuesta extractiva de preguntas, la clasificación de secuencias y la clasificación de tokens. Como alternativa, las tareas de generación de texto se pueden adaptar para ejecutarse en Inf1, según este tutorial de AWS Neuron MarianMT. Puedes encontrar más información sobre los modelos que están listos para usarse en Inferentia en la sección Model Architecture Fit de la documentación de Neuron.

Dependencias

Usar AWS Neuron para convertir modelos requiere las siguientes dependencias y entornos:

Un entorno Neuron SDK, que viene preconfigurado en AWS Deep Learning AMI.

Convertir un modelo a AWS Neuron

Con el mismo script usado en Uso de TorchScript en Python para rastrear un "BertModel", puedes importar la extensión del framework torch.neuron para acceder a los componentes del SDK de Neuron a través de una API de Python.

from transformers import BertModel, BertTokenizer, BertConfig
import torch
import torch.neuron

Y modificando la línea de código de rastreo de:

torch.jit.trace(model, [tokens_tensor, segments_tensors])

con lo siguiente:

torch.neuron.trace(model, [token_tensor, segments_tensors])

Este cambio permite a Neuron SDK rastrear el modelo y optimizarlo para ejecutarse en instancias Inf1.

Para obtener más información sobre las funciones, las herramientas, los tutoriales de ejemplo y las últimas actualizaciones de AWS Neuron SDK, consulte la documentación de AWS NeuronSDK.