企业官网建设流程全解析

如何为 Transformers Trainer 自定义 loss 函数？

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

Transformers 的Trainer默认直接使用模型 forward 返回的 loss（通常是交叉熵）。如果你的任务需要“loss 的计算方式”不同于默认逻辑——例如基于 logits 自己算 loss、引入参考模型、或改变归一化方式——文档给出了两条路径：把自定义函数通过compute_loss_func参数传给Trainer（轻量，只改 loss 计算，forward pass 不变），或者子类化并重写~Trainer.compute_loss（需要修改 forward pass 或 loss 计算的整体流程）。本文按这两条路径给出可照做的写法，依据是仓库中的 Trainer features 文档、Subclassing Trainer methods 文档 与 Gradient accumulation 文档 的 “Loss scaling” 一节。

先判断走哪条路径

trainer_customize.md 对两种定制方式做了明确区分：

• compute_loss_func在 forward pass之后运行，只定义“如何从模型输出计算 loss”，不触碰 forward pass；
• 需要修改 forward pass 本身（模型看不到 labels、输入被拼接、需要额外参考模型等）时，子类化~Trainer.compute_loss；
• 如果只是控制“何时发生”某些行为（日志、评估、早停），应该用 Callback，而不是改 loss。

判断标准：你的改动只影响“从输出到 loss 的公式”就走第一条；影响模型看到什么输入、返回什么就走第二条。

主路径：通过 compute_loss_func 替换默认 loss

1. 写一个符合签名的 loss 函数

文档要求的签名（来自 trainer_recipes.md）：

import torch.nn.functional as F def my_loss_fn(outputs, labels, num_items_in_batch): logits = outputs["logits"] loss = F.cross_entropy(logits, labels, reduction="sum") return loss / num_items_in_batch

三个参数的含义（均为文档原文定义）：

• outputs：模型原始输出，outputs.logits形状为(batch, seq_len, vocab_size)；
• labels：[Trainer] 在 forward pass之前从输入 batch 中弹出的 token ids；
• num_items_in_batch：整个累积 batch 中预测目标的数量。对 causal LM 模型，它统计的是 shifted labels（labels[..., 1:]），因为标签移位后每个序列的第 0 位没有目标。

关键点：提供自定义 loss 函数后，[Trainer] 会跳过自动的 loss 归一化，归一化必须由你的函数自己完成，上例用loss / num_items_in_batch实现按 token 数归一化。

2. 把函数传给 Trainer

trainer = Trainer( model=model, args=TrainingArguments(...), train_dataset=train_dataset, compute_loss_func=my_loss_fn, ) trainer.train()

其中TrainingArguments(...)处填入你自己训练的超参数（output_dir等），文档示例未展开这部分。

3. 与梯度累积配合时的 loss 缩放

按 grad_accumulation.md 的 “Loss scaling” 一节：

• 在自定义 loss 中使用num_items_in_batch，[Trainer] 会把 loss 除以所有 mini-batch 的预测目标总数，即按 token 归一化，而不是按gradient_accumulation_steps的固定步数；
• 如果不处理，[Trainer] 会改为把 loss 除以gradient_accumulation_steps；
• causal LM 场景下num_items_in_batch统计 shifted labels：预测位置i对应位置i + 1的 token，每个序列位置 0 没有目标，[Trainer] 会对labels[..., 1:]计数，保证分母与 loss 实际使用的预测目标数一致；
• 当数据 collator 直接提供shift_labels（如无 padding 的 collator）时，[Trainer] 改为统计该张量；masked LM 和分类等其他 loss 类型统计完整 label 张量。

替代路径：子类化 Trainer.compute_loss

当 loss 计算无法套进默认的“model 返回 loss”流程时——trainer_customize.md 以 DPO 为例说明：模型永远看不到 labels（只返回 logits）、chosen 与 rejected 响应被拼接、还需要一个参考模型各自的 log-probs——这时重写~Trainer.compute_loss。其签名为：

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 风格重写示例（展示整体结构，get_logps为示例中自定义的 log-prob 计算）：

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

两个使用限制（来自同一文档）：

• 完整可子类化的方法列表以 [Trainer] API 文档为准；
• 带下划线前缀的私有方法（如_save_checkpoint、_evaluate）也可以覆盖，但可能无预警变更，不要依赖。

验证与已记录的异常现象

文档没有给出“训练成功的固定日志”，但源码给出了两种可直接对号判断的现象（trainer.py）：

1. labels 缺失时的警告：提供了compute_loss_func但输入中没有labels时，[Trainer] 会打印警告Trainer:compute_loss_funcis defined butlabels=None. Your custom loss function will still be called with labels=None.，你的函数仍会被以labels=None调用——如果你的 loss 依赖 labels，说明数据侧的列名或 collator 有问题；
2. 归一化精度提示：compute_loss的 docstring 明确说明，如果你的 loss 不使用num_items_in_batch，需要把self.model_accepts_loss_kwargs覆盖为False，否则梯度累积时 loss 计算可能略有偏差。

按上面的签名写完函数、传入compute_loss_func后，运行trainer.train()，训练循环中每一步的 loss 即来自你的函数；若看到第 1 条警告或 loss 值与预期量级不符（未做归一化），按第 2 条检查参数传递方式。

参考文档

• 自定义 loss 函数主路径：docs/source/en/trainer_recipes.md（Custom loss function 一节）
• 子类化compute_loss及与 Callback 的分工：docs/source/en/trainer_customize.md
• 梯度累积下的 loss 缩放规则：docs/source/en/grad_accumulation.md（Loss scaling 一节）
• 实现细节：src/transformers/trainer.py

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers