PyTorch自定义损失函数教程

本文深入解析了PyTorch中自定义损失函数的核心实践与常见陷阱：强调虽非强制，但强烈建议继承`nn.Module`以保障参数管理、设备迁移、状态保存及训练流程集成；明确指出自动求导依赖张量运算的计算图完整性，而非是否继承模块，并系统梳理了导致梯度中断的典型错误——如混用NumPy/Python原生函数、显式`.item()`或`.numpy()`调用、不可微控制流（如`if`语句）以及未将loss压缩为标量；同时通过Focal Loss等实例详解正确实现要点，并提供实用调试技巧，包括验证梯度回传、手动提取梯度、监控梯度范数及规范超参类型管理，助你写出健壮、可维护、真正参与端到端训练的自定义损失函数。

自定义损失函数必须继承 nn.Module 吗？

不是必须，但强烈建议。直接写一个纯 Python 函数（比如 def my_loss(pred, target): return torch.mean((pred - target) ** 2)）也能用，但它无法参与 nn.Module 的参数管理、设备迁移（.to(device)）、状态保存等流程；更重要的是，它没法被当成模型的一部分嵌入训练循环（比如放进 model.loss_fn）。继承 nn.Module 是为了统一生命周期和行为，不是为了“能求导”——求导靠的是张量运算本身是否在计算图中。

为什么 forward 里不能用 NumPy 或 Python 内置 math？

因为 PyTorch 的自动微分只跟踪 torch.Tensor 上的运算，且要求所有操作是可导的、支持反向传播的算子。一旦混入 np.log、math.exp、list.append 或显式 .item() / .numpy() 调用，计算图就断了，loss.backward() 会报 RuntimeError: element 0 of tensors does not require grad 或更隐蔽的梯度为 None。

• ✅ 正确：torch.log(pred.clamp(min=1e-8))（使用 torch 版本 + 防零）
• ❌ 错误：np.log(pred.cpu().numpy())、math.log(pred.item())、pred.detach().numpy()
• ⚠️ 注意：torch.where 可导，但 if pred > 0.5: 这类控制流会破坏图，需改用 torch.where 或 torch.sigmoid 等向量化替代

怎么确保梯度能正确回传到模型输出？

核心就一条：所有中间变量必须由输入张量经可导运算得到，且最终 loss 是标量（shape == () 或 shape == (1,)）。常见翻车点：

• 忘了 .mean() 或 .sum() —— 如果返回的是 (N,) 形状的张量，backward() 会要求传入 gradient 参数，否则报错 grad can be implicitly created only for scalar outputs
• 用了 torch.nn.functional 里的函数但没注意其默认行为，比如 F.cross_entropy(input, target, reduction='none') 返回向量，而 reduction='mean' 才返回标量
• 手动实现了类似 Dice 系数的损失，但除法用了 / 之外还加了 + 1e-6，这个常数没问题，但若写成 denominator.item() + 1e-6 就断图了

示例（带权重的 Focal Loss 片段）：

class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
def forward(self, inputs, targets):
    # inputs: (N, C), targets: (N,)
    log_pt = F.log_softmax(inputs, dim=1)
    pt = torch.exp(log_pt)
    log_pt = log_pt.gather(1, targets.unsqueeze(1)).squeeze(1)  # (N,)
    pt = pt.gather(1, targets.unsqueeze(1)).squeeze(1)         # (N,)
    focal_weight = ((1 - pt) ** self.gamma)
    loss = -self.alpha * focal_weight * log_pt
    return loss.mean()  # ← 关键：必须压缩为标量

调试时怎么确认梯度真的传回去了？

别只看 loss.backward() 不报错——那只是说明当前这步没崩。真正要验证的是模型最后一层权重的 .grad 是否非空且合理：

• 在 loss.backward() 后立刻检查：print(model.fc.weight.grad is not None)
• 如果为 False，说明梯度没传到那里，大概率是某处用了 .detach()、.no_grad() 上下文，或 loss 没依赖该参数（比如 forward 里漏写了该层）
• 进阶技巧：用 torch.autograd.grad(loss, [model.fc.weight], retain_graph=True) 手动提取梯度，绕过 zero_grad() 干扰
• 可视化梯度分布：torch.norm(model.fc.weight.grad) 在训练初期应随 epoch 缓慢下降，突变为 0 或 inf 就得查

最常被忽略的一点：自定义 loss 类里定义的属性（如 self.alpha）如果是 Python float，它不会自动转成 torch.Tensor，但在运算中会被隐式提升；可读性与可控性更好的做法是用 nn.Parameter(torch.tensor(alpha), requires_grad=False) 显式声明——虽然它不参与优化，但统一了类型和设备管理。

今天带大家了解了的相关知识，希望对你有所帮助；关于文章的技术知识我们会一点点深入介绍，欢迎大家关注golang学习网公众号，一起学习编程~