| """ |
| 损失函数模块 — Person C 负责实现 |
| |
| 功能要求: |
| 1. LabelSmoothedCrossEntropyLoss: 带标签平滑的交叉熵损失 |
| |
| 技术要点: |
| - 标签平滑 (Label Smoothing) 是 Transformer 训练的标准技巧 |
| - smoothing=0.1 意味着将 10% 的概率均匀分配给非目标类 |
| - 需要忽略 padding token 的损失 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class LabelSmoothedCrossEntropyLoss(nn.Module): |
| """ |
| 带标签平滑的交叉熵损失。 |
| |
| 标签平滑将真实标签的概率质量从 1.0 重新分配: |
| target token 概率 = 1 - smoothing |
| 其他 token 概率 = smoothing / (V - 1) |
| |
| 使用 KL 散度实现: loss = KL(smooth_target || log_softmax(logits)) |
| |
| 参考: "Rethinking the Inception Architecture for Computer Vision" (Szegedy et al.) |
| """ |
|
|
| def __init__(self, smoothing: float = 0.1, pad_id: int = 0): |
| super().__init__() |
| self.smoothing = smoothing |
| self.pad_id = pad_id |
| self.confidence = 1.0 - smoothing |
|
|
| def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: |
| """ |
| Args: |
| logits: [B, T, V] — 模型输出 |
| targets: [B, T] — 目标 token ids (padding 位置应为 pad_id, 默认 -100) |
| Returns: |
| loss: scalar |
| """ |
| logits = logits.contiguous() |
| targets = targets.contiguous() |
|
|
| batch_size, seq_len, vocab_size = logits.size() |
|
|
| logits_flat = logits.view(-1, vocab_size) |
| targets_flat = targets.view(-1) |
|
|
| log_probs = F.log_softmax(logits_flat, dim=-1) |
|
|
| |
| non_pad_mask = targets_flat.ne(self.pad_id).float() |
|
|
| if self.smoothing > 0.0: |
| smooth_dist = torch.full_like(log_probs, self.smoothing / (vocab_size - 1)) |
| |
| safe_targets = targets_flat.clamp(min=0) |
| smooth_dist.scatter_(1, safe_targets.unsqueeze(1), self.confidence) |
| nll_loss = -torch.sum(smooth_dist * log_probs, dim=-1) |
| nll_loss = nll_loss * non_pad_mask |
| else: |
| nll_loss = F.nll_loss( |
| log_probs, targets_flat, |
| ignore_index=self.pad_id, reduction="none", |
| ) |
|
|
| num_tokens = non_pad_mask.sum().clamp(min=1) |
| loss = nll_loss.sum() / num_tokens |
|
|
| return loss |
|
|
