Transformer代码详解与项目实战之Masking

在此模型中Masking有两种，分别是Padding Masking和Look-ahead Masking。
Padding Masking：遮挡一批序列中所有的填充标记（pad tokens）。这确保了模型不会将填充作为输入。此模型使用0作为填充数据，mask的值表明填充值 0 出现的位置：在这些位置 mask 输出 1，否则输出 0。
Look-ahead Masking：前瞻遮挡（look-ahead mask）用于遮挡一个序列中的后续标记（future tokens）。换句话说，该 mask 表明了不应该使用的条目。这意味着要预测第三个词，将仅使用第一个和第二个词。与此类似，预测第四个词，仅使用第一个，第二个和第三个词，依此类推。

import tensorflow as tf

def masking(sequence, task='Padding'):
    """
    :param sequence: 输入tensor
    :param task: 分为"Padding"和"Sequence"(Look-ahead),默认为Padding
    :return:
    """

    if task == "Padding":
        zeroT = tf.cast(tf.math.equal(sequence, 0), tf.float32)   #元素为0的位置标记为1，其余位置标记为0
        return zeroT[:, tf.newaxis, tf.newaxis, :]  #构造四个维度，为了应该用时映射[batch_size, num_head, seq_len, seq_len]

    elif task == "Sequence":
        size = sequence.shape[1]
        triMatrix = tf.linalg.band_part(tf.ones((size, size)), -1, 0)  #只保留对角线及下三角阵内容
        return 1 - triMatrix  #上三角阵（不包括对角线）内容为1，代表Attention时候看不到的内容

    else:
        raise ValueError("任务名称只能是“Padding”或“Sequence”")

def make_masking(inp, tar):
    enc_padding_mask = masking(inp, task='Padding')
    dec_padding_mask = masking(inp, task='Padding')

    look_ahead_mask = masking(tar, task='Sequence')
    dec_tar_padding_mask = masking(tar, task='Padding')

    combined_mask = tf.maximum(look_ahead_mask, dec_tar_padding_mask)  #取对应元素较大值。numpy的broadcast原理

    return enc_padding_mask, combined_mask, dec_padding_mask

make_masking为模块函数入口，返回值分别是encoder中的masking，decoder第一层的masking和第二层的masking。