两同架构CNN训练时mask网络梯度全为0，求问题原因

我有两个架构相同但初始化方式不同的CNN网络，分别叫net和mask。我的训练循环代码如下：

for epoch in range(epochs):

    print(f"Epoch {epoch + 1}")

    net.train()
    mask.train()

    # 用mask剪枝net
    net_masked = apply_mask(net, mask)
    
    running_loss = 0.0

    for images, labels in train_loader:
      images, labels = images.to(device), labels.to(device)

      # 用图像计算net的损失
      outputs_net = net_masked(images)
      loss_net = criterion(outputs_net, labels)

      outputs_mask = mask(images) # 同时计算mask的损失（没用到loss_mask.item()）
      loss_mask = criterion(outputs_mask, labels)

      loss = loss_net + loss_mask*0 # 现在loss是这样的：tensor(2.2969, grad_fn=<AddBackward0>)

      optimizer_mask.zero_grad()

      loss.backward() # 尝试计算梯度

      optimizer_mask.step()  # 更新权重

      running_loss += loss_net.item()

    train_loss.append(running_loss / len(train_loader))

我想要更新mask的权重，以此每次用不同方式剪枝net，但计算梯度时发现所有梯度都是0。请问为什么无法计算梯度？是损失计算方式有问题吗？

梯度全为0的核心问题完全出在损失计算逻辑上：

1. 损失与mask权重无关联 #

你写的loss = loss_net + loss_mask*0，相当于最终损失就是loss_net，而loss_net是经过apply_mask(net, mask)得到的输出计算的损失，但这里存在两个致命问题：

• 如果apply_mask是不可导的硬剪枝操作（比如直接把权重置0），那么loss_net的梯度无法反向传播到mask的权重上——剪枝操作直接切断了梯度流。
• 就算apply_mask是可导的软掩码操作（比如用乘法实现），你把loss_mask*0直接清零，相当于完全移除了mask的训练信号，mask的权重没有任何损失来驱动更新，自然梯度为0。

2. 正确的损失设计方向 #

要让mask获得有效梯度，必须让最终损失和mask的输出/权重直接挂钩：

• 监督式剪枝：保留mask的分类损失项，赋予合理权重，让mask学会输出有效的分类结果，同时间接完成剪枝目标：

  loss = loss_net + 0.1 * loss_mask  # 0.1为超参数，可根据任务调整
  

• 稀疏性约束：加入稀疏性损失（比如L1正则），鼓励mask输出更多0，实现剪枝的稀疏性目标：

  import torch
  sparse_loss = torch.norm(torch.cat([p.flatten() for p in mask.parameters()]), p=1)
  loss = loss_net + 0.01 * sparse_loss  # 0.01为正则权重
  

3. 额外关键检查 #

• 确认apply_mask是可导操作：如果是硬剪枝（直接置0），必须改成软掩码（比如用sigmoid将mask输出限制在0-1之间，再与net的权重相乘），否则梯度无法流向mask。
• 确认optimizer_mask仅针对mask的参数初始化：optimizer_mask = torch.optim.Adam(mask.parameters(), lr=1e-3)，避免误包含net的参数。
• 若仅训练mask，建议将net设为eval()模式或冻结其参数，防止net的权重被意外更新。

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