PyTorch MLP训练异常：损失恒定且权重未更新问题排查

我搭建了一个用于轨迹二分类的DCR网络，其架构与训练代码如下：

网络架构 #

class DCR(nn.Module):
    def __init__(self, kemb_size, nvar, points, device):
        super().__init__()
        self.phis = load_phis_dataset()
        self.kemb = get_kernel_embedding(self.phis, nvar, samples = kemb_size).to(device) # (concepts, kemb_size)
        _ = self.kemb.requires_grad_()
        
        self.fc1 = nn.Linear(kemb_size + (nvar*points), 64)
        self.fc2 = nn.Linear(64, 1)
        self.sigmoid = nn.Sigmoid()
                
    def forward(self, x):
        # concept truth degrees
        rhos = get_robustness(x, self.phis, time = False) # (trajectories, concepts)
        _ = rhos.requires_grad_()
        # embed trajectories in kernel space
        traj_emb = torch.matmul(rhos, self.kemb) # (trajectories, kemb_size)
        _ = traj_emb.requires_grad_()
        # combine info from traj_embed and x to predict class 
        x_new = x.view(x.size(0), -1) # flatten x
        combined_features = torch.cat((traj_emb, x_new), dim=1) # (trajectories, kemb_size + x.shape[0]*x.shape[1])

        output = self.fc1(combined_features)
        output = F.relu(output)
        output = self.fc2(output)
        output = self.sigmoid(output)
        
        return output.squeeze(1)

训练代码 #

model = DCR(kemb_size, nvar, points, device).to(device)
criterion = nn.BCELoss().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01)

model.train()
for epoch in range(10):
    epoch_loss = 0.0

    for batch, labels in train_loader:
        batch, labels = batch.to(device), labels.to(device)
                 
        y_preds = model(batch)
        loss = criterion(y_preds, labels.float())
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        epoch_loss += y_preds.shape[0] * loss.item()
  
    print(f'Epoch: {epoch}, Loss: {epoch_loss/len(y_train):.5f}')

然而训练时损失在每个epoch都完全恒定，且权重未更新。我尝试过设置.requires_grad_()但没有效果，请问可能是什么问题导致的？

补充说明：

• phis = STL公式列表
• kemb = 上述STL公式的核嵌入
• rhos = 输入轨迹上STL公式的鲁棒性
  所有张量形状均如代码注释所示。

以下是几个核心原因及对应解决方案：

1. 计算图断开：自定义函数未支持自动微分 #

• 问题：get_robustness和get_kernel_embedding如果使用了numpy或其他非PyTorch张量操作，会导致rhos或kemb成为无梯度传递的叶子张量，即使手动调用.requires_grad_()，上游计算也没构建有效计算图，梯度无法传到全连接层。
• 验证：在loss.backward()后打印model.fc1.weight.grad，若为None则说明梯度未传递；再检查rhos.grad和kemb.grad是否为None，定位断开点。
• 解决：确保两个自定义函数的所有计算都用PyTorch张量实现；若必须用外部代码，需用torch.autograd.Function封装成可微分操作。

2. kemb未被注册为可训练参数 #

• 问题：你手动设置了kemb.requires_grad_()，但model.parameters()只会遍历PyTorch默认注册的参数（如nn.Linear的权重）。如果kemb是直接赋值的普通张量而非nn.Parameter，它不会被加入优化器的参数集合，自然不会被更新。
• 解决：将kemb定义为nn.Parameter，自动注册为模型可训练参数：

  self.kemb = nn.Parameter(get_kernel_embedding(self.phis, nvar, samples = kemb_size).to(device))
  

  无需再手动调用.requires_grad_()。

3. 优化器配置不合理 #

• 问题：当前SGD优化器的学习率0.01可能过低，导致权重更新幅度极小，看起来像是没有变化；或者SGD本身不适合当前任务的收敛特性。
• 解决：尝试调高学习率至0.1/0.5测试，或更换为自适应学习率的Adam优化器：

  optimizer = optim.Adam(model.parameters(), lr=0.001)
  

4. 损失函数输入饱和导致梯度消失 #

• 问题：如果y_preds输出全集中在0或1附近，Sigmoid函数在该区间的梯度趋近于0，BCELoss的梯度会几乎为0，导致权重无法更新。
• 解决：去掉最后一层的Sigmoid，改用BCEWithLogitsLoss（合并Sigmoid与BCELoss，数值稳定性更好，避免梯度消失）：

  # 网络forward中移除sigmoid
  def forward(self, x):
      # ... 现有代码 ...
      output = self.fc2(output)
      return output.squeeze(1)
  
  # 训练代码更换损失函数
  criterion = nn.BCEWithLogitsLoss().to(device)
  

5. 数据或标签异常 #

• 问题：若所有batch的标签全为0或1，模型无法学习有效特征；或者train_loader未开启shuffle=True，每个epoch输入完全一致，也会出现损失恒定的假象。
• 验证：打印几个batch的labels检查分布，确认DataLoader是否设置shuffle=True。
• 解决：确保数据集标签分布均衡，且DataLoader开启shuffle=True。

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