ForwardKLLoss¶
- class torchtune.modules.loss.ForwardKLLoss(ignore_index: int = - 100)[源代码]¶
用于有效索引的 Kullback-Leibler 散度损失。的实现https://github.com/jongwooko/distillm/blob/17c0f98bc263b1861a02d5df578c84aea652ee65/distillm/losses.py
- 参数:
ignore_index (int) – 指定一个被忽略且不会对输入梯度做出贡献的目标值。损失在非忽略目标上进行划分。默认值:-100。
- forward(student_logits: Tensor , teacher_logits: Tensor , labels: Tensor ) Tensor [源代码]¶
- 参数:
student_logits (torch.Tensor) – 来自学生模型的 logits,形状为 (batch_size*num_tokens, vocab_size)。
teacher_logits (torch.Tensor) – 来自教师模型的 logits,形状为 (batch_size*num_tokens, vocab_size)。
labels (torch.Tensor) – 形状为 (batch_size, vocab_size) 的地面实况标签。
- 返回值:
形状为 (1,) 的 KL 散度损失。
- 返回类型: