在此模型中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。