Trainer API

1. Por qué existe el Trainer

En el curso de PyTorch aprendiste a escribir el bucle de entrenamiento manualmente: forward pass, calcular loss, backward pass, actualizar pesos. Ese bucle tiene exactamente 10-15 líneas de Python. Pero cuando pasas a modelos grandes y entrenamiento serio, ese bucle simple se expande enormemente:

Bucle básico que conoces (PyTorch):
  for batch in dataloader:
      outputs = model(batch)
      loss = criterion(outputs, labels)
      loss.backward()
      optimizer.step()
      optimizer.zero_grad()

Bucle "de producción" (lo que habría que agregar):
  ✗ Mover datos a GPU (batch.to(device))
  ✗ Evaluación periódica en validation set
  ✗ Guardar checkpoint del mejor modelo
  ✗ Logging de métricas (loss, accuracy, tiempo)
  ✗ Gradient clipping (evitar gradientes explosivos)
  ✗ Learning rate scheduling (warmup + decay)
  ✗ Gradient accumulation (para GPUs pequeñas)
  ✗ Mixed precision training (FP16/BF16)
  ✗ EarlyStopping (parar si no mejora)
  ✗ Distribución en múltiples GPUs
  ✗ Métricas personalizadas por época

→ Son +100 líneas adicionales de código repetitivo.
→ El Trainer encapsula todo esto en 5-10 líneas de configuración.

El Trainer no hace nada mágico que no puedas hacer tú mismo — solo agrupa esas 100 líneas en un objeto configurable. Si en algún momento necesitas más control, puedes subclasificarlo y sobreescribir métodos específicos. Por ahora, lo usamos como viene.

2. TrainingArguments: todos los controles del entrenamiento

TrainingArguments es un objeto de configuración. Lo creas antes del Trainer y le dices exactamente cómo quieres que entrene: cuántas épocas, qué tan rápido aprender, cuándo guardar, etc. Vamos parámetro por parámetro, sin saltar ninguno.

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./mi-modelo-finetuned",
    num_train_epochs=4,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    warmup_steps=100,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="eval_accuracy",
    logging_steps=50,
    report_to="none",
)

Ahora explicamos cada parámetro:

output_dir — dónde guardar todo

La carpeta donde el Trainer guardará los checkpoints (copias del modelo en distintos momentos), los logs y el modelo final. Si no existe, la crea. Usa una ruta descriptiva que recuerde el experimento.

num_train_epochs — cuántas veces ver los datos

Una época (en inglés, epoch) es una pasada completa por todos los datos de entrenamiento. Si tienes 8.000 ejemplos y batch_size=16, una época = 8.000 ÷ 16 = 500 pasos de actualización de pesos.

Para clasificación de texto (modelos tipo BERT):
  · 2-5 épocas suele ser suficiente.
  · Más de 5 épocas → riesgo de sobreajuste (memorizar en vez de aprender).
  · Señal de alerta: val_loss sube mientras train_loss sigue bajando.

Para fine-tuning de modelos generativos (LLaMA, Mistral):
  · 1-3 épocas suele ser suficiente.
  · Estos modelos aprenden rápido y se sobreajustan fácilmente.

Para datasets pequeños (< 1.000 ejemplos):
  · Hasta 10 épocas puede ser razonable.
  · Siempre monitorear val_loss.

per_device_train_batch_size y per_device_eval_batch_size

Cuántos ejemplos procesar juntos en cada paso. "Per device" significa "por GPU/CPU". Si tienes 2 GPUs y per_device_train_batch_size=16, el batch total efectivo es 32. Para evaluación puedes usar un número mayor porque no se calculan gradientes, así que ocupa menos memoria.

learning_rate — el tamaño del paso

Este es el número que multiplica el gradiente antes de restar de los pesos. Lo estudiaste en el curso de gradiente. Para fine-tuning de modelos preentrenados se usan valores muy pequeños: entre 1e-5 y 5e-5. Si es demasiado grande, el modelo "olvida" lo que aprendió antes (catastrophic forgetting). Si es demasiado pequeño, aprende muy lento.

3. warmup_steps — empezar despacio y acelerar

El warmup es una técnica donde el learning rate empieza en un valor muy pequeño y va creciendo gradualmente hasta llegar al valor que configuraste. Después de los pasos de warmup, el learning rate empieza a decrecer suavemente hasta llegar a cero al final del entrenamiento.

Configuración: learning_rate=2e-5, warmup_steps=100, total_steps=500

Paso   0:  lr = 0.00000002 × (0/100)  =  0.0000000  (empieza en cero)
Paso  20:  lr = 0.00000002 × (20/100) =  0.000000004 (20% del lr máximo)
Paso  50:  lr = 0.00000002 × (50/100) =  0.000000010 (50% del lr máximo)
Paso 100:  lr = 0.00000002 × (100/100)=  0.000000020 (100% — lr máximo)
   ↑ aquí termina el warmup, empieza el decay ↑
Paso 200:  lr ≈ 0.000000015 (75% — bajando gradualmente)
Paso 300:  lr ≈ 0.000000010 (50%)
Paso 400:  lr ≈ 0.000000005 (25%)
Paso 500:  lr ≈ 0.0000000   (0% — el último paso)

¿Cuántos warmup_steps usar?
  → Regla práctica: 5-10% del total de pasos
  → Con 1000 ejemplos, batch=16, 3 épocas: total_steps = (1000/16) × 3 ≈ 188 pasos
  → warmup_steps = 188 × 0.06 ≈ 11 pasos

Cuando hay pocos datos (menos de 1.000 ejemplos), puedes poner warmup_steps=0 y no pasa nada. El warmup es más importante cuando el dataset es grande y el modelo es sensible.

4. weight_decay — frenando los pesos grandes

El weight decay (también llamado regularización L2) es una penalización que se añade al loss para evitar que los pesos del modelo se vuelvan muy grandes. Lo viste en el curso de redes neuronales como "L2 regularization".

En el curso de redes, la fórmula del loss con L2 era:

Donde λ (lambda, la letra griega que se lee "lambda") es el peso de la regularización. En HuggingFace, weight_decay=0.01 significa λ = 0.01. En práctica, esto hace que en cada actualización los pesos se "encojan" ligeramente hacia cero, lo que evita el sobreajuste.

Sin weight_decay:
  w_nuevo = w_actual - lr × gradiente
  w_nuevo = 0.5 - 0.00002 × 3.0 = 0.5 - 0.00006 = 0.49994

Con weight_decay=0.01:
  w_nuevo = w_actual × (1 - lr × weight_decay) - lr × gradiente
  factor  = 1 - 0.00002 × 0.01 = 1 - 0.0000002 = 0.9999998
  w_nuevo = 0.5 × 0.9999998 - 0.00002 × 3.0
          = 0.4999999 - 0.00006 = 0.49994

El efecto es minúsculo en cada paso, pero acumulado en miles de pasos:
  · Pesos que no son útiles tienden a cero (el gradiente es pequeño).
  · Pesos importantes se mantienen grandes (el gradiente los empuja).
  → El modelo aprende a ser selectivo en qué pesos importan.

5. El ciclo interno del Trainer — paso a paso

Cuando llamas a trainer.train(), el Trainer ejecuta este ciclo. Lo escribimos explícitamente para que no haya nada oculto:

El Trainer también maneja automáticamente:

• gradient clipping — si un gradiente se dispara a un valor enorme (a veces pasa con texto largo), el Trainer lo recorta a un valor máximo para que el entrenamiento no explote.
• mixed precision — si tu GPU lo soporta y tú lo activas con fp16=True, usa FP16 para los cálculos y ahorra memoria.
• guardar checkpoints — al final de cada época (o cada N pasos), guarda el estado completo del modelo para que puedas retomar si algo falla.

6. Callbacks: interrumpir el entrenamiento cuando ya no mejora

Un callback (literalmente "llamada de vuelta") es una función que el Trainer llama en momentos específicos: al iniciar el entrenamiento, al final de cada paso, al final de cada época, etc. Los callbacks te permiten agregar comportamiento personalizado sin tocar el bucle principal.

EarlyStoppingCallback — el más importante

El EarlyStoppingCallback monitorea la métrica de validación y para el entrenamiento automáticamente si esa métrica no mejora durante N evaluaciones consecutivas. Esto evita que el modelo siga entrenando (y sobreajustando) cuando ya encontró su mejor punto.

Métrica monitoreada: eval_loss (buscamos que BAJE)
patience=2 (aguanta 2 épocas sin mejora antes de parar)

Época 1: eval_loss = 0.450  → nuevo mínimo, guardar modelo
Época 2: eval_loss = 0.380  → nuevo mínimo, guardar modelo
Época 3: eval_loss = 0.362  → nuevo mínimo, guardar modelo
Época 4: eval_loss = 0.371  → subió. Contador de paciencia: 1/2
Época 5: eval_loss = 0.375  → sigue subiendo. Contador: 2/2
          → PARAR. El mejor modelo fue el de la época 3.
          (Si patience=3, habría esperado una época más)

from transformers import (
    TrainingArguments, Trainer,
    EarlyStoppingCallback,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
)

# Los argumentos de entrenamiento, con early stopping activado
args = TrainingArguments(
    output_dir="./modelo-finetuned",
    num_train_epochs=10,               # ponemos 10, pero EarlyStopping parará antes
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    warmup_steps=100,
    weight_decay=0.01,
    evaluation_strategy="epoch",     # OBLIGATORIO para early stopping
    save_strategy="epoch",           # guardar en cada época para poder cargar el mejor
    load_best_model_at_end=True,     # al parar, cargar el mejor checkpoint
    metric_for_best_model="eval_loss", # qué métrica mirar (queremos que BAJE)
    greater_is_better=False,          # False porque buscamos que el loss BAJE
    logging_steps=50,
    report_to="none",
)

# Crear el Trainer con el callback de early stopping
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=ds_tok["train"],
    eval_dataset=ds_tok["validation"],
    tokenizer=tokenizer,
    data_collator=collator,
    compute_metrics=compute_metrics,
    callbacks=[
        EarlyStoppingCallback(
            early_stopping_patience=2,   # aguanta 2 épocas sin mejora
            early_stopping_threshold=0.001  # mejora mínima para no contar como estancado
        )
    ],
)

trainer.train()

Otros callbacks disponibles

HuggingFace incluye más callbacks útiles:

7. Ejemplo completo funcional — de principio a fin

Juntamos todo: cargar modelo, preparar datos (los del Lección 6), configurar el Trainer con todos los parámetros explicados y entrenarlo. Este código funciona en Google Colab con GPU gratuita.

import torch
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
    EarlyStoppingCallback,
)
import evaluate

## ── 1. Dataset ────────────────────────────────────────
raw = load_dataset("imdb")
# Usamos solo un subconjunto para ir más rápido en el ejemplo
raw["train"] = raw["train"].select(range(5000))
raw["test"]  = raw["test"].select(range(1000))

# Crear validation split (10% del train)
splits = raw["train"].train_test_split(test_size=0.1, seed=42)
raw["train"]      = splits["train"]
raw["validation"] = splits["test"]

## ── 2. Tokenizar ──────────────────────────────────────
MODEL_NAME = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

def tokenize_fn(batch):
    return tokenizer(batch["text"], truncation=True, max_length=256)

ds = raw.map(tokenize_fn, batched=True)
collator = DataCollatorWithPadding(tokenizer=tokenizer)

## ── 3. Modelo ─────────────────────────────────────────
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)

## ── 4. Métricas ───────────────────────────────────────
metric = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    return metric.compute(predictions=preds, references=labels)

## ── 5. TrainingArguments ──────────────────────────────
args = TrainingArguments(
    output_dir="./imdb-distilbert",
    num_train_epochs=8,                  # early stopping parará antes
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    warmup_steps=100,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="eval_accuracy",
    greater_is_better=True,
    logging_steps=50,
    report_to="none",
)

## ── 6. Trainer ────────────────────────────────────────
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=ds["train"],
    eval_dataset=ds["validation"],
    tokenizer=tokenizer,
    data_collator=collator,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
)

trainer.train()

Tres cosas importantes para notar en esta salida:

1. El train loss sigue bajando de 0.45 a 0.08 — el modelo parece mejorar.
2. El eval loss empieza a subir desde la época 4 — el modelo está empezando a memorizar.
3. El EarlyStopping detectó que la accuracy no mejoró desde la época 3 (2 épocas seguidas) y paró — cargando el mejor modelo.

🎮 Configurador de hiperparámetros

Ajusta los sliders y ve cómo cambia el comportamiento esperado del entrenamiento. La demo te da retroalimentación sobre si la configuración es recomendable o no.

9. Lo que aprendiste

En la próxima lección: Evaluación y métricas — por qué accuracy sola miente, cómo leer la matriz de confusión, qué es F1 y cómo saber si tu modelo realmente funciona.