Callbacks

TrainerCallback se integra en eventos específicos del entrenamiento (inicio de época, evaluación, finalización del entrenamiento) para modificar el estado del entrenamiento o el flujo de control. Utilice callbacks para registrar métricas en herramientas de seguimiento de experimentos como Trackio, personalizar cuándo se realizan el guardado y la evaluación, o añadir otro comportamiento personalizado. Apile varios callbacks para combinar funcionalidades.

Los callbacks no pueden modificar el propio bucle de entrenamiento, como la pasada hacia adelante. Para cambiar lo que Trainer calcula, subclase sus métodos.

El diagrama siguiente muestra todos los eventos en los que un callback puede engancharse.

on_train_begin
  └─ for each epoch:
      on_epoch_begin
      └─ for each step:
          on_step_begin
          ├─ for each gradient accumulation substep:
          │   on_substep_end
          on_pre_optimizer_step   ← after gradient clipping, before optimizer.step()
          on_optimizer_step       ← after optimizer.step()
          on_step_end
          └─ (conditionally):
              on_log
              on_evaluate
              on_save
      on_epoch_end
on_train_end
on_predict / on_prediction_step
on_push_begin

Creación de un callback

Subclase TrainerCallback y sobrescribe uno o varios métodos de eventos del diagrama anterior. El siguiente ejemplo demuestra tres hooks:

• on_epoch_begin imprime el epoch actual desde TrainerState (un valor en vivo) y la tasa de aprendizaje desde TrainingArguments (un valor estático).
• on_step_end lee la tasa de aprendizaje actual desde lr_scheduler (pasada a través de **kwargs) y establece should_evaluate en TrainerControl (consulte TrainerControl para obtener una lista completa de objetos de control) para desencadenar una evaluación cuando la tasa desciende por debajo del umbral.
• on_train_end imprime la mejor métrica y el paso en el que se produjo.

from transformers import TrainerCallback

class EpochLoggerCallback(TrainerCallback):
    def __init__(self, lr_eval_threshold=1e-5):
        self.lr_eval_threshold = lr_eval_threshold

    def on_epoch_begin(self, args, state, control, **kwargs):
        print(f"Starting epoch {int(state.epoch) + 1}/{state.num_train_epochs} "
              f"(lr={args.learning_rate})")

    def on_step_end(self, args, state, control, **kwargs):
        lr_scheduler = kwargs.get("lr_scheduler")
        if lr_scheduler is not None:
            current_lr = lr_scheduler.get_last_lr()[0]
            if current_lr < self.lr_eval_threshold:
                control.should_evaluate = True

    def on_train_end(self, args, state, control, **kwargs):
        print(f"Training complete! Best metric: {state.best_metric} at step {state.best_global_step}")

Registra el callback en Trainer mediante el argumento callbacks. También puedes pasar varios callbacks como una lista.

trainer = Trainer(
    callbacks=[EpochLoggerCallback(lr_eval_threshold=1e-5)],
    ...,
)

Callbacks integrados

Transformers incluye varios callbacks integrados que están activos de forma predeterminada. Los callbacks integrados adicionales registran información en plataformas como Trackio.

DefaultFlowCallback

DefaultFlowCallback manages the default logging, evaluation, and checkpoint schedule based on the logging_strategy, eval_strategy, and save_strategy values in TrainingArguments. At the right step or epoch, it sets the corresponding control flags (should_log, should_evaluate, should_save). It also sets should_training_stop when global_step reaches max_steps.

La forma principal de personalizar cuándo se realiza el registro, la evaluación o el guardado es sobrescribir esta devolución de llamada.

ProgressCallback y PrinterCallback

Trainer elige automáticamente entre estas dos callbacks en función del campo disable_tqdm de TrainingArguments.

• ProgressCallback se utiliza de forma predeterminada. Muestra una barra de progreso de tqdm durante el entrenamiento y una barra separada durante la evaluación o la predicción, e imprime las métricas más recientes en cada evento on_log. Durante el entrenamiento distribuido, solo se ejecuta en el proceso principal para evitar salidas duplicadas.
• PrinterCallback se utiliza cuando disable_tqdm=True. Imprime el diccionario de registro en stdout en cada evento on_log sin barra de progreso.

También puedes intercambiarlos manualmente con remove_callback() y add_callback().

from transformers import PrinterCallback

trainer = Trainer(...)
trainer.remove_callback(ProgressCallback)
trainer.add_callback(PrinterCallback)

EarlyStoppingCallback

EarlyStoppingCallback detiene el entrenamiento cuando una métrica de evaluación deja de mejorar. Después de cada evaluación, comprueba si la métrica mejoró en más de early_stopping_threshold. Si la métrica no ha mejorado durante early_stopping_patience evaluaciones consecutivas, el entrenamiento se detiene.

EarlyStoppingCallback requiere dos TrainingArguments:

• metric_for_best_model, la métrica de evaluación a supervisar
• eval_strategy, si se evalúa en "steps" o en "epoch"

from transformers import EarlyStoppingCallback

trainer = Trainer(
    ...,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.001)],
)

Próximos pasos

• Vea todos los callbacks integrados disponibles para registrar en los rastreadores de experimentos.
• La guía Métodos de subclase de Trainer explica cómo sobrescribir los métodos de Trainer cuando necesitas modificar lo que calcula el bucle de entrenamiento.