Pytorch实战Transformer算法之Masks

Transformer 中，遮罩（masking）的概念非常重要，有两种 masks，分别是 padding mask 和 look ahead mask。

1：Padding mask

在训练的时候，序列的长度是不一样的，因此需要对较短的序列进行填充（pad）以使它们的长度相同。这些填充的部分是不需要“关注”的，因为它们不包含实际的信息。

def create_padding_mask(seq):
    # padding mask 就是将索引序列中为 0 的位置设为 1，其余位置设为 0
    mask = (seq == 0).float()
    # broadcasting，在第一维和第二维添加新维度
    return mask.unsqueeze(1).unsqueeze(2)

vocab_size = 2008
emsize = 200#d_model
inp = torch.LongTensor([[1,2,4,5,0],[0,6,7,8,9]]) 
tar = torch.LongTensor([[2,4,5,0,0,6],[6,7,8,9,10,0]])

inp_mask = create_padding_mask(inp)
tar_mask = create_padding_mask(tar)

print("mask维度：",inp_mask.shape)
print("mask：",inp_mask)
print(torch.squeeze(inp_mask),torch.squeeze(inp_mask).shape)

mask维度： torch.Size([2, 1, 1, 5])
mask： tensor([[[[0., 0., 0., 0., 1.]]],


        [[[1., 0., 0., 0., 0.]]]])
tensor([[0., 0., 0., 0., 1.],
        [1., 0., 0., 0., 0.]]) torch.Size([2, 5])

可以看到，无需“关注的位置”，Mask 是 1。

2：Look ahead mask

Look-ahead mask 主要用于 Decoder，确保在自注意力机制中，每个 Token 只会关注到前面的 Token，而不会提前关注到“未来”的 Token。这样可以防止模型在生成序列时记住不该记住的内容。

def create_look_ahead_mask(size):
    """Create a look-ahead mask with a right upper triangular shape.

    Args:
        size: The size of the mask (size x size).

    Returns:
        mask: A tensor of shape (size, size) with a right upper triangular shape.
    """
    # 创建一个全1的矩阵
    mask = torch.ones((size, size))
    # 使用 torch.triu 创建一个右上角的三角形遮罩，对角线和对角线右上方的元素为 1，其余为 0
    mask = torch.triu(mask, diagonal=1)
    return mask

测试一个序列的 look_ahead_mask：

seq_len = emb_inp.shape[1]   
look_ahead_mask = create_look_ahead_mask(seq_len)
print("emb_inp:", emb_inp)
print("look_ahead_mask：", look_ahead_mask,look_ahead_mask.shape) # torch.Size([5, 5])

look_ahead_mask： tensor([[0., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1.],
        [0., 0., 0., 1., 1.],
        [0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0.]]) torch.Size([5, 5])

是不是一个很典型的右上角矩阵，和 Padding mask 一样，位置为 1 的内容将来需要忽略。

来源：aigcrepo