pytorch中如何给网络添加mask

深度学习中，我们经常会遇到需要添加mask的场景，如：

• nlp中为了长度对齐，需要补齐长度，但在计算attention时会将补齐位置mask掉从而不参与attention计算；
• mask相关的预训练任务，如MLM、MAE等，需要mask掉被遮盖的token，以完成预测的预训练任务；
• swin中，在做shift操作后，为了防止原本物理位置不相邻的区域产生交互，需要进行mask attention。
• 计算loss时想忽略掉一些不想用来计算该loss的样本。

样例

在attention操作中，在计算attn softmax前，将被mask位置的logits设置为一个很小的数，如-10000，在计算softmax后，就会抑制掉这些位置的作用，代码如下：

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., with_qkv=True):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5  # 分母根号d
        self.with_qkv = with_qkv
        if self.with_qkv:
           self.qkv = nn.Linear(dim, dim * 3, bia