在 PyTorch 中训练期间使用掩码更改元素是否有效？

Question

我有一个正向函数，它包含一个线性回归任务，看起来像这样

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[indicator_arr == False] = x_hat #Here is the questionable part
        ...
        ...
        "different operations which include the variable x only"
        ...

forward

函数返回一个值，该值取决于

x

并用于损失函数。

indicator_arr

是一个包含

True

或

False

值的数组。目的是屏蔽矩阵

x

的特定部分，并用

x``<sub>hat</sub>

改变它们以进行进一步处理。

self.bbox_model

是一个线性层，其中只训练偏置向量。

出现的问题是偏置向量的梯度非常低，实际上接近于零，因此没有进行训练，我的验证损失很高。掩蔽是否是造成这种情况的可能原因，或者在训练期间此掩蔽步骤是否会发生任何其他错误？

我希望反向传播函数可以正确地回溯所有内容，并且优化器会更新偏置向量。

Answer 1

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):
        mask = torch.tensor(indicator_arr, dtype=torch.float32, device=x.device)
        x_hat = self.bbox_model(data)
        masked_x_hat = mask * x_hat
        masked_x = mask * x
        masked_x += masked_x_hat
        ...
        ...
        "different operations which include the variable masked_x only"
        ...