Callbacks

TrainerCallback se integra a eventos específicos do treinamento (início de epoch, avaliação, fim do treinamento) para modificar o estado do treinamento ou o fluxo de controle. Use callbacks para registrar métricas em rastreadores de experimentos como Trackio, personalizar quando o salvamento e a avaliação ocorrem, ou adicionar outros comportamentos personalizados. Empilhe múltiplos callbacks para combinar funcionalidades.

Os callbacks não podem modificar o próprio loop de treinamento, como a passagem direta. Para alterar o que o Trainer calcula, crie uma subclasse dos seus métodos.

O diagrama abaixo mostra todos os eventos em que um callback pode se conectar.

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

Criando um callback

Subclasse TrainerCallback e sobrescreva um ou mais métodos de eventos do diagrama acima. O exemplo abaixo demonstra três hooks:

• on_epoch_begin imprime o epoch atual de TrainerState (um valor dinâmico) e a taxa de aprendizado de TrainingArguments (um valor estático).
• on_step_end lê a taxa de aprendizado atual do lr_scheduler (passado via **kwargs) e define should_evaluate em TrainerControl (veja TrainerControl para uma lista completa dos objetos de controle) para disparar uma avaliação quando ela cair abaixo do limite.
• on_train_end imprime a melhor métrica e o passo em que ela ocorreu.

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

Registre o callback no Trainer no argumento callbacks. Você também pode passar vários callbacks como uma lista.

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

Callbacks integrados

O Transformers inclui vários callbacks integrados que estão ativos por padrão. Callbacks integrados adicionais, como callbacks integrados, registram logs em plataformas como Trackio.

DefaultFlowCallback

DefaultFlowCallback gerencia o agendamento padrão de registro de logs, avaliação e checkpoint com base nos valores de logging_strategy, eval_strategy e save_strategy em TrainingArguments. No passo ou epoch adequado, ele define as flags correspondentes de control (should_log, should_evaluate, should_save). Ele também define should_training_stop quando global_step atinge max_steps.

A sobrescrita deste callback é a principal forma de personalizar quando o registro, a avaliação ou o salvamento ocorrem.

ProgressCallback e PrinterCallback

O Trainer escolhe automaticamente entre esses dois callbacks com base no campo disable_tqdm em TrainingArguments.

• ProgressCallback é usado por padrão. Ele exibe uma barra de progresso tqdm durante o treinamento e uma barra separada durante a avaliação ou a predição, e imprime as métricas mais recentes a cada evento on_log. Durante o treinamento distribuído, ele é executado apenas no processo principal para evitar saída duplicada.
• PrinterCallback é usado quando disable_tqdm=True. Ele imprime o dicionário de log no stdout a cada evento on_log, sem barra de progresso.

Você também pode trocá-los manualmente com remove_callback() e add_callback().

from transformers import PrinterCallback

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

EarlyStoppingCallback

EarlyStoppingCallback interrompe o treinamento quando uma métrica de avaliação deixa de melhorar. Após cada avaliação, ele verifica se a métrica melhorou em mais de early_stopping_threshold. Se a métrica não tiver melhorado em early_stopping_patience avaliações consecutivas, o treinamento é interrompido.

EarlyStoppingCallback requer dois TrainingArguments:

• metric_for_best_model, a métrica de avaliação a ser monitorada
• eval_strategy, se a avaliação deve ser feita em "steps" ou "epoch"

from transformers import EarlyStoppingCallback

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

Próximas etapas

• Veja todos os callbacks integrados disponíveis para registro em rastreadores de experimentos.
• O guia Subclassing Trainer methods aborda a sobrescrita dos métodos de Trainer quando você precisa alterar o que o loop de treinamento calcula.