コールバック

TrainerCallback は、特定のトレーニングイベント（エポック開始、評価、トレーニング終了）にフックし、トレーニング状態の変更や制御フローの操作を行います。コールバックを使用して、Trackio などの実験トラッカーへのメトリクスの記録、保存や評価の実行タイミングのカスタマイズ、その他のカスタム動作の追加を行うことができます。複数のコールバックをスタックすることで、機能を組み合わせることができます。

Callbackは、フォワードパスのような学習ループ自体を変更することはできません。Trainerが計算する内容を変更するには、代わりにそのメソッドをサブクラス化してください。

以下の図は、コールバックがフックできるすべてのイベントを示しています。

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 をサブクラス化し、上記の図にあるイベントメソッドのうち1つ以上をオーバーライドします。以下の例では、3つのフックを示しています。

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

callbacks 引数で、Trainer にコールバックを登録します。リストとして複数のコールバックを渡すこともできます。

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

組み込みコールバック

Transformers には、デフォルトで有効になる組み込みのコールバックがいくつか含まれています。追加の統合コールバックは、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 フィールドに基づいて、これらの2つのコールバックのいずれかを自動的に選択します。

ProgressCallbackはデフォルトで使用されます。トレーニング中はtqdmのプログレスバーを表示し、評価または予測時には別のバーを表示し、各on_logイベントで最新のメトリクスを出力します。分散学習時には、重複した出力を避けるため、メインプロセスでのみ実行されます。
PrinterCallback は、disable_tqdm=True の場合に使用されます。進捗バーを表示せず、on_log イベントのたびにログ辞書を標準出力に出力します。

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が2つ必要です:

監視する評価指標である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)],
)

次のステップ

実験トラッカーへのログ出力に対応するすべての統合コールバックを確認できます。
Trainer メソッドのサブクラス化ガイドでは、トレーニングループが計算する内容を変更する必要がある場合に Trainer のメソッドをオーバーライドする方法について説明しています。

Update on GitHub