深度学习中,我们经常会遇到需要添加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