# 子类化 Trainer 方法

通过子类化 [Trainer](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer) 的方法来更改训练行为，而无需重写整个循环。子类化会修改 *training loop*，例如前向传播或损失计算。

在创建子类之前，请先考虑你需要改变的是 [Trainer](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer) 计算*什么*，还是它*何时*以及*是否*执行操作。对于时序和条件逻辑，请改用 [Callback](./trainer_callbacks)。Callback 控制事件发生的时机（如日志记录、评估、早停），而子类化则改变事件发生的内容（如损失计算、数据加载、优化）。

> [!NOTE]
> 请参阅 [Trainer](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer) API 文档，以获取可子类化的完整方法列表。私有方法（以 `_` 为前缀）例如 `_save_checkpoint` 或 `_evaluate` 也可以被重写，但这些方法可能会在不另行通知的情况下发生变更。

## get_train_dataloader

标准的 [get_train_dataloader()](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer.get_train_dataloader) 方法会加载一个批次，在其上进行训练，然后丢弃该批次，再加载下一个批次。

```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 个 token 的补全需要约 512 次顺序前向传播，而一个训练步骤只需一次前向传播。[GRPOTrainer](https://huggingface.co/docs/trl/v1.17.1/en/grpo_trainer#trl.GRPOTrainer) 通过重写 [get_train_dataloader()](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer.get_train_dataloader) 来跨多个步骤对生成进行批处理。

`trl.GRPOTrainer.get_train_dataloader` 通过将 batch size 乘以 `steps_per_generation` 参数，一次性加载用于多个训练步骤的 *batches* 生成提示。如果 `train_batch_size=4` 且 `steps_per_generation=8`，dataloader 将生成大小为 32 的批次，从而将生成成本降低 8 倍。

```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/v5.23.0/zh/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/v1.17.1/en/dpo_trainer#trl.DPOTrainer) 重写了 [compute_loss()](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer.compute_loss)，因为损失计算在多个方面与标准交叉熵不同：

- 模型从未看到标签；它仅返回 logits，供 DPO 从中计算对数概率
- 选中的回复和被拒绝的回复会被拼接在一起
- 参考模型计算其自身的对数概率
- 损失是 `π_chosen`、`π_rejected`、`π_ref_chosen` 和 `π_ref_rejected` 的函数

以上均不符合标准的 [Trainer.compute_loss()](/docs/transformers/v5.23.0/zh/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/v1.17.1/en/grpo_trainer#trl.GRPOTrainer) 和 [DPOTrainer](https://huggingface.co/docs/trl/v1.17.1/en/dpo_trainer#trl.DPOTrainer) 如何扩展 [Trainer](/docs/transformers/v5.23.0/zh/main_classes/trainer#transformers.Trainer)，或 [Axolotl](https://github.com/axolotl-ai-cloud/axolotl/tree/main/src/axolotl/core/trainers) 如何在其基础上构建自定义训练器。
- 如果只需要自定义训练事件（例如在训练步骤结束时记录指标）期间发生的内容，请查阅 [Callbacks](./trainer_callbacks) 指南。