PyTorch如何正确实现带掩码的均值计算，避免反向传播报错

报错原因 #

报错的核心是rnn_out[text == self.pad_token] *= 0属于原地修改操作，直接修改了RNN输出的带梯度的张量，破坏了反向传播需要的计算图结构，即便套了torch.no_grad()也无法规避这个问题。

简洁实现方案 #

无需传入序列长度、无需遍历batch，直接通过广播掩码实现向量化计算，代码如下：

def forward(self, text):
    # 'text' 维度为 (B, L)，B为batch size，L为序列长度
    embed = self.embed(text)
    rnn_out, _ = self.rnn(embed)  # 输出维度为 (B, L, 2*H)，H为单方向隐层维度
    
    # 构造非pad位置的掩码，扩展维度适配广播规则
    mask = (text != self.pad_token).unsqueeze(-1)  # 维度变为 (B, L, 1)
    # 非原地操作置0pad对应位置的输出，不破坏梯度计算链
    rnn_out_masked = rnn_out * mask
    # 计算有效长度，加clamp避免极端情况出现除以0错误
    valid_len = mask.sum(dim=1).clamp(min=1)
    # 按序列维度求和后除以有效长度得到平均特征
    feat = rnn_out_masked.sum(dim=1) / valid_len
    
    feat = self.dropout(feat)
    return feat

方案优势 #

• 全向量化操作，运算效率远高于逐样本遍历实现
• 无任何原地修改操作，完全适配PyTorch自动梯度计算规则
• 自带鲁棒性处理，避免样本全为pad时的除0异常

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