# Trainerメソッドのサブクラス化

[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer) のメソッドをサブクラス化することで、ループ全体を書き直すことなく学習の動作を変更できます。サブクラス化により、*training loop*（例えばフォワードパスや損失の計算）が変更されます。

サブクラス化を検討する前に、[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)が*何を*計算するか、あるいは*いつ*および*どの条件で*動作するかを変更する必要があるかを検討してください。タイミングや条件付きロジックを制御する場合は、代わりに[Callback](./trainer_callbacks)を使用してください。Callbackはイベントの発生タイミング（ログ出力、評価、早期停止など）を制御し、サブクラス化は実行される処理の内容（損失の計算、データ読み込み、最適化など）を変更します。

> [!NOTE]
> サブクラス化できるメソッドの完全なリストについては、[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer) の API ドキュメントを参照してください。`_` で始まるプライベートメソッド（`_save_checkpoint` や `_evaluate` など）もオーバーライドできますが、これらは予告なく変更される場合があります。

## get_train_dataloader

標準的な[get_train_dataloader()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.get_train_dataloader)メソッドは、1つのバッチを読み込み、それに対して学習を行い、破棄して、次のバッチを読み込みます。

```py
def get_train_dataloader(self):
    return self._get_dataloader(
        batch_size=self._train_batch_size,
        ...
)
```

[GRPO](https://huggingface.co/docs/trl/en/grpo_trainer) は、学習前に補完を生成するオンライン強化学習アルゴリズムです。各ステップで補完を生成するのは、自己回帰的であるためコストが高くなります。512トークンの補完には、学習ステップの1回のフォワードパスに対して約512回の逐次的なフォワードパスが必要です。[GRPOTrainer](https://huggingface.co/docs/trl/main/en/grpo_trainer#trl.GRPOTrainer) は [get_train_dataloader()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.get_train_dataloader) をサブクラス化し、複数ステップにわたって生成をバッチ処理します。

`trl.GRPOTrainer.get_train_dataloader` は、バッチサイズに `steps_per_generation` 引数を乗算することで、複数の学習ステップ分の *バッチ* を一度に読み込みます。`train_batch_size=4` かつ `steps_per_generation=8` の場合、dataloader は 32 のバッチを生成し、生成コストを 8 分の 1 に削減します。

```py
def get_train_dataloader(self):
    dataloader_params = {
        "batch_size": self._train_batch_size * self.args.steps_per_generation, # this is the only change
        ...
    }
```

## 損失を計算

[compute_loss()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.compute_loss) は、モデルによって計算されたクロスエントロピー損失を返します。

```py
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    ...
    outputs = model(**inputs)
    ...
    loss = outputs["loss"] # get loss from model

    return (loss, outputs) if return_outputs else loss
```

[DPO](https://huggingface.co/docs/trl/en/dpo_trainer)は、参照モデルに対する相対的な観点から、ポリシーモデルが拒否された応答よりも選択された応答をどの程度強く好むかを測定します。[DPOTrainer](https://huggingface.co/docs/trl/main/en/dpo_trainer#trl.DPOTrainer)は、損失の計算が標準的なクロスエントロピーと複数の点で異なるため、[compute_loss()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.compute_loss)をサブクラス化しています。

- モデルはラベルを一切参照せず、DPOが対数確率を計算するためのlogitsのみを返します
- 選択された応答と拒否された応答が連結されます
- 参照モデルが自身のlog-probsを計算します
- 損失は`π_chosen`、`π_rejected`、`π_ref_chosen`、`π_ref_rejected`の関数です

上記のいずれも、標準的な[Trainer.compute_loss()](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer.compute_loss)メソッドには当てはまりません。

```py
def compute_loss(
    self,
    model: PreTrainedModel | nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    return_outputs=False,
    num_items_in_batch=None,
) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]:
    ...
    outputs = model(**inputs)
    logits = outputs.logits
    logps = get_logps(logits, inputs)
    chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    ref_logits = self.ref_model(**inputs).logits
    ref_logps = get_logps(ref_logits, inputs)
    ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    chosen_scores = chosen_logps - ref_chosen_logps
    rejected_scores = rejected_logps - ref_rejected_logps
    per_sequence_loss = -F.logsigmoid(self.beta * chosen_scores - rejected_scores)
    loss = per_sequence_loss.mean()
    return (loss, outputs) if return_outputs else loss
```

## 次のステップ

- より実用的な例として、TRL における [GRPOTrainer](https://huggingface.co/docs/trl/main/en/grpo_trainer#trl.GRPOTrainer) と [DPOTrainer](https://huggingface.co/docs/trl/main/en/dpo_trainer#trl.DPOTrainer) が [Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer) をどのように拡張しているか、または [Axolotl](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/core/trainers) がその上にカスタムトレーナーをどのように構築しているかを参照してください。
- トレーニングステップの終了時にメトリクスを記録するなど、トレーニングイベント中に発生する処理のみをカスタマイズする必要がある場合は、[Callbacks](./trainer_callbacks)ガイドを確認してください。