PyTorch训练中掩码修改元素是否合法？梯度近零问题排查

问题场景 #

我实现了一个仅训练偏置向量的线性回归模型，代码如下：

import torch
import torch.nn as nn

class lin_reg(nn.Module):
    def __init__(self, pre_trained_f):
        super(lin_reg, self).__init__()
        self.bbox_model = nn.Linear(4, 4, bias=True)
        # 固定权重为单位矩阵，仅训练偏置
        self.bbox_model.weight = torch.nn.Parameter(torch.eye(4))
        self.bbox_model.weight.requires_grad = False
        self.bbox_model.bias.requires_grad = True

    def forward(self, x, data, indicator_arr):
        x_hat = self.bbox_model(data)
        # 此处为疑似问题代码：用x_hat替换x中indicator_arr为False的部分
        x[indicator_arr == False] = x_hat
        # 后续仅基于x进行一系列运算，最终返回用于损失计算的值
        ...

模型的forward流程中，用indicator_arr作为掩码，将x的特定部分替换为bbox_model的输出x_hat，再参与后续损失计算。但训练时发现bbox_model的偏置向量梯度极低（接近0），模型完全没有训练进展，验证损失始终居高不下。

问题根源：原地修改破坏计算图追踪 #

你代码里的x[indicator_arr == False] = x_hat是原地（in-place）修改操作，这是导致梯度无法正常回传的核心原因：

• PyTorch的反向传播依赖计算图追踪张量的生成路径，原地操作会直接修改原始张量的存储，破坏了计算图中x_hat到后续运算的链路
• 当你原地修改x后，后续运算使用的x已经不是原来的计算图节点，x_hat对应的梯度无法正确回溯到bbox_model.bias

修复方案：避免原地操作，生成新张量 #

替换掉原地修改的代码，用torch.where生成新的张量，保留完整的计算图：

def forward(self, x, data, indicator_arr):
    x_hat = self.bbox_model(data)
    # 生成掩码，注意保证mask和x、x_hat的设备/数据类型一致
    mask = indicator_arr == False
    # 用torch.where生成新张量：mask为True的位置取x_hat，否则取x
    x_new = torch.where(mask, x_hat, x)
    # 后续所有运算改用x_new
    ...

额外检查点 #

1. 设备与数据类型一致性：确保indicator_arr、x、x_hat在同一设备（CPU/GPU），且indicator_arr是bool类型，避免隐式转换打断梯度传递
2. 其他原地操作排查：检查后续代码中是否还有类似x.add_(...)、x.copy_(...)的原地操作，这类操作同样会破坏计算图
3. 优化器参数配置：确认优化器确实将bbox_model.bias加入了待更新参数组，例如：

   optimizer = torch.optim.SGD([model.bbox_model.bias], lr=0.01)
   

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