Ajuste de un modelo con la API Trainer

🤗 Transformers incluye una clase Trainer para ayudarte a ajustar cualquiera de los modelos preentrenados proporcionados en tu dataset. Una vez que hayas hecho todo el trabajo de preprocesamiento de datos de la última sección, sólo te quedan unos pocos pasos para definir el Trainer. La parte más difícil será preparar el entorno para ejecutar Trainer.train(), ya que se ejecutará muy lentamente en una CPU. Si no tienes una GPU preparada, puedes acceder a GPUs o TPUs gratuitas en Google Colab.

Los siguientes ejemplos de código suponen que ya has ejecutado los ejemplos de la sección anterior. Aquí tienes un breve resumen de lo que necesitas:

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding

raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)


def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)


tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

Entrenamiento

El primer paso antes de que podamos definir nuestro Trainer es definir una clase TrainingArguments que contendrá todos los hiperparámetros que el Trainer utilizará para el entrenamiento y la evaluación del modelo. El único argumento que tienes que proporcionar es el directorio donde se guardarán tanto el modelo entrenado como los puntos de control (checkpoints). Para los demás parámetros puedes dejar los valores por defecto, deberían funcionar bastante bien para un ajuste básico.

from transformers import TrainingArguments

training_args = TrainingArguments("test-trainer")

El segundo paso es definir nuestro modelo. Como en el capítulo anterior, utilizaremos la clase AutoModelForSequenceClassification, con dos etiquetas:

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

Observarás que, a diferencia del Capítulo 2, aparece una advertencia después de instanciar este modelo preentrenado. Esto se debe a que BERT no ha sido preentrenado para la clasificación de pares de frases, por lo que la cabeza del modelo preentrenado se ha eliminado y en su lugar se ha añadido una nueva cabeza adecuada para la clasificación de secuencias. Las advertencias indican que algunos pesos no se han utilizado (los correspondientes a la cabeza de preentrenamiento eliminada) y que otros se han inicializado aleatoriamente (los correspondientes a la nueva cabeza). La advertencia concluye animándote a entrenar el modelo, que es exactamente lo que vamos a hacer ahora.

Una vez que tenemos nuestro modelo, podemos definir un Trainer pasándole todos los objetos construidos hasta ahora: el model, los training_args, los datasets de entrenamiento y validación, nuestro data_collator, y nuestro tokenizer:

from transformers import Trainer

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

Ten en cuenta que cuando pasas el tokenizer como hicimos aquí, el data_collator por defecto utilizado por el Trainer será un DataCollatorWithPadding como definimos anteriormente, por lo que puedes omitir la línea data_collator=data_collator. De todas formas, era importante mostrarte esta parte del proceso en la sección 2.

Para ajustar el modelo en nuestro dataset, sólo tenemos que llamar al método train() de nuestro Trainer:

trainer.train()

Esto iniciará el ajuste (que debería tardar un par de minutos en una GPU) e informará de la training loss cada 500 pasos. Sin embargo, no te dirá lo bien (o mal) que está rindiendo tu modelo. Esto se debe a que:

1. No le hemos dicho al Trainer que evalúe el modelo durante el entrenamiento especificando un valor para evaluation_strategy: steps (evaluar cada eval_steps) o epoch (evaluar al final de cada época).
2. No hemos proporcionado al Trainer una función compute_metrics() para calcular una métrica durante dicha evaluación (de lo contrario, la evaluación sólo habría impreso la pérdida, que no es un número muy intuitivo).

Evaluación

Veamos cómo podemos construir una buena función compute_metrics() para utilizarla la próxima vez que entrenemos. La función debe tomar un objeto EvalPrediction (que es una tupla nombrada con un campo predictions y un campo label_ids) y devolverá un diccionario que asigna cadenas a flotantes (las cadenas son los nombres de las métricas devueltas, y los flotantes sus valores). Para obtener algunas predicciones de nuestro modelo, podemos utilizar el comando Trainer.predict():

predictions = trainer.predict(tokenized_datasets["validation"])
print(predictions.predictions.shape, predictions.label_ids.shape)

(408, 2) (408,)

La salida del método predict() es otra tupla con tres campos: predictions, label_ids, y metrics. El campo metrics sólo contendrá la pérdida en el dataset proporcionado, así como algunas métricas de tiempo (cuánto se tardó en predecir, en total y de media). Una vez que completemos nuestra función compute_metrics() y la pasemos al Trainer, ese campo también contendrá las métricas devueltas por compute_metrics().

Como puedes ver, predictions es una matriz bidimensional con forma 408 x 2 (408 es el número de elementos del dataset que hemos utilizado). Esos son los logits de cada elemento del dataset que proporcionamos a predict() (como viste en el capítulo anterior, todos los modelos Transformer devuelven logits). Para convertirlos en predicciones que podamos comparar con nuestras etiquetas, necesitamos tomar el índice con el valor máximo en el segundo eje:

import numpy as np

preds = np.argmax(predictions.predictions, axis=-1)

Ahora podemos comparar esas predicciones preds con las etiquetas. Para construir nuestra función compute_metric(), nos basaremos en las métricas de la librería 🤗 Evaluate. Podemos cargar las métricas asociadas al dataset MRPC tan fácilmente como cargamos el dataset, esta vez con la función evaluate.load(). El objeto devuelto tiene un método compute() que podemos utilizar para calcular de la métrica:

import evaluate

metric = evaluate.load("glue", "mrpc")
metric.compute(predictions=preds, references=predictions.label_ids)

{'accuracy': 0.8578431372549019, 'f1': 0.8996539792387542}

Los resultados exactos que obtengas pueden variar, ya que la inicialización aleatoria de la cabeza del modelo podría cambiar las métricas obtenidas. Aquí, podemos ver que nuestro modelo tiene una precisión del 85,78% en el conjunto de validación y una puntuación F1 de 89,97. Estas son las dos métricas utilizadas para evaluar los resultados en el dataset MRPC para la prueba GLUE. La tabla del paper de BERT recoge una puntuación F1 de 88,9 para el modelo base. Se trataba del modelo “uncased” (el texto se reescribe en minúsculas antes de la tokenización), mientras que nosotros hemos utilizado el modelo “cased” (el texto se tokeniza sin reescribir), lo que explica el mejor resultado.

Juntándolo todo obtenemos nuestra función compute_metrics():

def compute_metrics(eval_preds):
    metric = evaluate.load("glue", "mrpc")
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

Y para ver cómo se utiliza para informar de las métricas al final de cada época, así es como definimos un nuevo Trainer con nuestra función compute_metrics():

training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

Ten en cuenta que hemos creado un nuevo TrainingArguments con su evaluation_strategy configurado como "epoch" y un nuevo modelo. De lo contrario sólo estaríamos continuando el entrenamiento del modelo que ya habíamos entrenado. Para lanzar una nueva ejecución de entrenamiento, ejecutamos:

trainer.train()

Esta vez, nos informará de la pérdida de validación y las métricas al final de cada época, además de la pérdida de entrenamiento. De nuevo, la puntuación exacta de precisión/F1 que alcances puede ser un poco diferente de la que encontramos nosotros, debido a la inicialización aleatoria del modelo, pero debería estar en el mismo rango.

El Trainer funciona en múltiples GPUs o TPUs y proporciona muchas opciones, como el entrenamiento de precisión mixta (usa fp16 = True en tus argumentos de entrenamiento). Repasaremos todo lo que ofrece en el capítulo 10.

Con esto concluye la introducción al ajuste utilizando la API de Trainer. En el Capítulo 7 se dará un ejemplo de cómo hacer esto para las tareas más comunes de PLN, pero ahora veamos cómo hacer lo mismo en PyTorch puro.