콜백

콜백 (Callbacks)

TrainerCallback은 특정 학습 이벤트(에폭 시작, 평가, 학습 종료)에 연결되어 학습 상태를 수정하거나 제어 흐름을 바꿔요. 콜백을 사용해 메트릭을 Trackio 같은 실험 추적기에 로깅하고, 저장과 평가가 일어나는 시점을 커스터마이즈하며, 다른 커스텀 동작을 추가할 수 있어요. 기능을 결합하려면 여러 콜백을 쌓아서 사용해요.

출처: 문서

본문

콜백은 forward pass 같은 학습 루프 자체를 수정할 수는 없어요. Trainer가 계산하는 것을 바꾸려면 subclass로 그 메서드를 재정의해야 해요.

아래 다이어그램은 콜백이 연결될 수 있는 모든 이벤트를 보여줘요.

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

콜백 만들기

TrainerCallback을 서브클래싱하고 위 다이어그램의 이벤트 메서드 중 하나 이상을 재정의해요. 아래 예시는 세 가지 훅을 보여줘요:

  • on_epoch_begin: TrainerState의 현재 에폭(실시간 값)과 TrainingArguments의 학습률(정적 값)을 출력해요.
  • on_step_end: lr_scheduler(**kwargs로 전달됨)에서 현재 학습률을 읽고, 임계값 아래로 떨어지면 TrainerControl의 should_evaluate를 설정해서 평가를 트리거해요 (컨트롤 객체 전체 목록은 TrainerControl 참고).
  • on_train_end: 최고 메트릭과 그것이 발생한 스텝을 출력해요.
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}")

Trainer에 callbacks 인자로 콜백을 등록해요. 여러 콜백을 리스트로 전달할 수도 있어요.

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

내장 콜백 (Built-in callbacks)

Transformers에는 기본적으로 활성화되는 여러 내장 콜백이 있어요. 추가적인 통합 콜백(integrated callbacks)은 Trackio 같은 플랫폼에 로깅해요.

DefaultFlowCallback

DefaultFlowCallback은 TrainingArguments의 logging_strategy, eval_strategy, save_strategy 값에 따라 기본 로깅, 평가, 체크포인트 일정을 관리해요. 올바른 스텝이나 에폭에서 대응하는 control 플래그(should_log, should_evaluate, should_save)를 설정해요. 또한 global_step이 max_steps에 도달하면 should_training_stop도 설정해요.

이 콜백을 재정의하는 것이 로깅, 평가, 저장이 언제 일어나는지를 커스터마이즈하는 주요 방법이에요.

ProgressCallback와 PrinterCallback

Trainer는 TrainingArguments의 disable_tqdm 필드에 따라 이 두 콜백 중 하나를 자동으로 선택해요.

  • ProgressCallback은 기본으로 사용돼요. 학습 중에 tqdm 진행 막대를 표시하고, 평가나 예측 중에는 별도의 막대를 표시하며, 각 on_log 이벤트에서 최신 메트릭을 출력해요. 분산 학습 중에는 중복 출력을 피하기 위해 메인 프로세스에서만 실행돼요.
  • PrinterCallback은 disable_tqdm=True일 때 사용돼요. 진행 막대 없이 매 on_log 이벤트에서 로그 딕셔너리를 stdout에 출력해요.

remove_callback()과 add_callback()으로 수동으로 교체할 수도 있어요.

from transformers import PrinterCallback

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

EarlyStoppingCallback

EarlyStoppingCallback은 평가 메트릭이 개선되지 않으면 학습을 중지해요. 각 평가 후에 메트릭이 early_stopping_threshold보다 많이 개선되었는지 확인해요. 메트릭이 early_stopping_patience번 연속으로 개선되지 않으면 학습이 중지돼요.

EarlyStoppingCallback은 두 개의 TrainingArguments가 필요해요:

  • metric_for_best_model: 모니터링할 평가 메트릭
  • eval_strategy: "steps" 또는 "epoch" 중 어느 쪽에서 평가할지
from transformers import EarlyStoppingCallback

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

다음 단계