如何为 Transformers Trainer 自定义 loss 函数? 📅 发布时间:2026/9/11 5:36:24 👁 浏览次数: 如何为 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/transformersTransformers 的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 替换默认 loss1. 写一个符合签名的 loss 函数文档要求的签名来自 trainer_recipes.mdimport torch.nn.functional as F def my_loss_fn(outputs, labels, num_items_in_batch): logits outputs[logits] loss F.cross_entropy(logits, labels, reductionsum) return loss / num_items_in_batch三个参数的含义均为文档原文定义outputs模型原始输出outputs.logits形状为(batch, seq_len, vocab_size)labels[Trainer] 在 forward pass之前从输入 batch 中弹出的 token idsnum_items_in_batch整个累积 batch 中预测目标的数量。对 causal LM 模型它统计的是 shifted labelslabels[..., 1:]因为标签移位后每个序列的第 0 位没有目标。关键点提供自定义 loss 函数后[Trainer] 会跳过自动的 loss 归一化归一化必须由你的函数自己完成上例用loss / num_items_in_batch实现按 token 数归一化。2. 把函数传给 Trainertrainer Trainer( modelmodel, argsTrainingArguments(...), train_datasettrain_dataset, compute_loss_funcmy_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_stepscausal 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_outputsFalse, num_items_in_batchNone): ... 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_outputsFalse, num_items_in_batchNone, ) - 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, dim0) # 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, dim0) 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.pylabels 缺失时的警告提供了compute_loss_func但输入中没有labels时[Trainer] 会打印警告Trainer:compute_loss_funcis defined butlabelsNone. Your custom loss function will still be called with labelsNone.你的函数仍会被以labelsNone调用——如果你的 loss 依赖 labels说明数据侧的列名或 collator 有问题归一化精度提示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.mdCustom loss function 一节子类化compute_loss及与 Callback 的分工docs/source/en/trainer_customize.md梯度累积下的 loss 缩放规则docs/source/en/grad_accumulation.mdLoss 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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考