自定义Cross Entropy Loss损失值与官方一致但模型不收敛求助

核心问题分析 #

虽然你的损失函数计算的数值结果和官方一致，但梯度传播路径异常是模型无法收敛的根本原因，以下是具体问题点：

• 梯度传播断裂：

  1. 初始loss是手动创建的CPU张量（torch.FloatTensor([0])），若模型在GPU上运行，这个张量会和模型参数不在同一设备，导致梯度无法反向传播到模型。
  2. 循环中每次对loss重新赋值（loss = loss - ...）会破坏计算图的连续性，PyTorch对这种循环累加的梯度追踪存在缺陷，容易导致梯度无法正确传递。

• 数值稳定性隐患：
  先执行softmax再取log的操作会放大数值下溢风险（比如softmax输出接近0时，log会变成-inf），虽然你加了torch.max(..., -100)截断，但官方CrossEntropyLoss是用log_softmax合并操作，从数学上避免了这种误差，同时梯度计算更精准。

• 冗余的Variable使用：
  PyTorch 0.4+版本已废弃手动Variable包装，直接设置张量的requires_grad=True即可，手动包装反而可能引入计算图异常。

优化后的实现代码 #

用批量操作替代循环，保证计算图连续，同时对齐官方实现的数值稳定性：

class MyLossFunc(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, out, label):
        # 用log_softmax替代先softmax再log，提升数值稳定性
        log_probs = torch.nn.functional.log_softmax(out, dim=1)
        # 批量提取对应类别的log概率，避免循环
        selected_log_probs = log_probs[range(len(label)), label]
        # 计算平均损失，和官方CrossEntropyLoss一致（官方默认reduction='mean'）
        loss = -torch.mean(selected_log_probs)
        return loss

额外验证步骤 #

1. 检查设备一致性：确保所有张量（模型参数、输入out、label、损失张量）都在同一设备（CPU/GPU）上，可以用out.device获取设备，创建张量时指定device=out.device。
2. 验证梯度：在训练前，手动计算模型参数的梯度，对比自定义损失和官方损失的梯度是否一致，比如：

   model.zero_grad()
   custom_loss = MyLossFunc()(out, label)
   custom_loss.backward()
   custom_grads = [p.grad.clone() for p in model.parameters()]
   
   model.zero_grad()
   official_loss = nn.CrossEntropyLoss()(out, label)
   official_loss.backward()
   official_grads = [p.grad.clone() for p in model.parameters()]
   
   # 检查梯度是否近似相等
   for c_grad, o_grad in zip(custom_grads, official_grads):
       print(torch.allclose(c_grad, o_grad))
   

内容的提问来源于stack exchange，提问作者zhuzhu