Transformer源码中Mask机制的实现
训练过程中的 Mask实现
mask 机制的原理是, 在 decoder 端, 做 self-Attention 的时候, 不能 Attention 还未被预测的单词, 预测的信息是基于encoder 与以及预测出的单词. ⽽在 encoder 阶段的, Self_Attention 却没有这个机制, 因为encoder 的self-Attention 是对句⼦中的所有单词 Attention ,mask 本质是对于 Attention 来说的, 所以我们来看下 Attention 的实现:
def attention(query, key, value, mask=None, dropout=None):
"Compute 'Scaled Dot Product Attention'"
d_k = query.size(-1)
scores = torch.matmul(query, anspose(-2, -1))  / math.sqrt(d_k)
# 这⾥是对应公式的  Q* K的转秩矩阵
"""
Queries张量，形状为[B, H, L_q, D_q]
Keys张量，形状为[B, H, L_k, D_k]
Values张量，形状为[B, H, L_v, D_v]，⼀般来说就是k
"""
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
p_attn = F.softmax(scores, dim = -1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
我们知道, 在训练的时候, 我们是以 batch_size 为单位的, 那么就会有 padding, ⼀般我们取 pad == 0,
那么就会造成在 Attention 的时候, query 的值为 0, query 的值为 0, 所以我们计算的对应的 scores 的值也是 0, 那么就会导致 softmax 很可能分配给该单词⼀个相对不是很⼩的⽐例，因此, 我们将 pad 对应的 score 取值为负⽆穷, 以此来减⼩ pad 的影响. 也就是上⾯中scores = scores.masked_fill(mask == 0, -1e9)的意思. 需要注意的是, 这⼀步在 Query 与 Key 相乘之后, 在 Softmax 之前, 所以本质上, decoder 的Self-Attention 的内部, Target 句⼦, 是会和 Encoder-Output 做 Self-Attention 的乘积的, 但是, 这个结果 Attention 的结果, 我们只取前⾯⼀部分, 取⼀部分的⽅法就是这⾥的mask. 我们可以⽤下⾯这张图来表⽰:
也就是说, 在 decoder, 未预测的单词也是⽤ padding 的⽅式加⼊到 batch 的(注意 Padding 其实就是通过mask 01矩阵实现的), 所以使⽤的mask 机制与 padding 时mask 的机制是相同的, 本质上都是query在做Attention之后的值为设为最⼩值, 只是 mask 矩阵不同, 如上图所⽰, decoder 在先⾃⾝做完 self_Attention 之后, 未被预测的单词的权重是最⼩值, 那么输出就是 0.
我们可以根据 decoder forward部分的代码发现这⼀点.
class DecoderLayer(nn.Module):
"Decoder is made of self-attn, src-attn, and feed forward (defined below)"
def forward(self, x, memory, src_mask, tgt_mask):
"Follow Figure 1 (right) for connections."
m = memory
# ⽬标语⾔的self_Attention, 这⾥ mask的作⽤就是⽤到上⾯所说的 softmax 之前的部分
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
# 这⾥使⽤的是 Self-Attention 机制，其实 m 是encoder的输出,x是decoder第⼀部分的输出,
# 因为上⾯⼀部分的输出中, 未被预测的单词的 query 其实是 0(padding), 那么在这⾥可以直接使⽤ src_mask
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
# 最后是两个线形层,
return self.sublayer[2](x, self.feed_forward)
接下来我们来追溯⼀下, 这⾥的 mask 是怎么来的, 我们最终构建的模块是 Encoder_Decoder,
class EncoderDecoder(nn.Module):
"""
A standard Encoder-Decoder architecture. Base for this and many
other models.
"""
def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
super(EncoderDecoder, self).__init__()
self.decoder = decoder
self.src_embed = src_embed
# 将源语⾔的单词 embedding 放在⼀起, position embedding
<_embed = tgt_embed
# 将⽬标语⾔的单词 embedding 放在⼀起, position embedding
# 就是最后产⽣结果的地⽅
def forward(self, src, tgt, src_mask, tgt_mask):
"Take in and process masked src and target sequences."
return self.de(src, src_mask), src_mask, tgt, tgt_mask)
def encode(self, src, src_mask):
der(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
return self._embed(tgt), memory, src_mask, tgt_mask)
我们在训练的时候, 使⽤的是 model.forward, 这⼀部分在:
def run_epoch(args, data_iter, model, loss_compute, valid_params=None, epoch_num=0,
is_valid=False, is_test=False, logger=None):
"Standard Training and Logging Function"
start = time.time()
total_tokens = 0
total_loss = 0
tokens = 0
if valid_params is not None:
src_dict, tgt_dict, valid_iter = valid_params
hist_valid_scores = []
bleu_all = 0
count_all = 0
for i, batch in enumerate(data_iter):
out = model.forward(batch.src, ,batch.src_mask, _mask)
# 参数来⾃ batch
loss = loss_compute(out, _y, kens)
# 这⼀步既计算了损失, ⼜更新了参数
total_loss += loss
total_tokens += kens
tokens += kens
这些都是训练的步骤, 数据是怎么来的, mask 矩阵来⾃ batch, 所以最关键的是 batch 是怎么来的, 再往回在 train.py函数中, 我们发现
_, logger_file = train_utils.run_epoch(args, (batch(pad_idx, b) for b in train_iter),
model_parallel if args.multi_gpu else model, train_loss_fn,
valid_params=valid_params,
epoch_num=epoch, logger=logger_file)
batch 是来⾃ rebatch 函数, 以及训练数据的迭代器, 这个train_iter 是根据 torchtext 得到, 这⾥就不赘述了, 所以关键就是下⾯的 rebatch 函数, def rebatch(pad_idx, batch):
"Fix order in torchtext"
src, trg = anspose(0, 1), anspose(0, 1)
# 读的数据是 sequence * batch_size 维度, 是在torchtext 中的Filed 决定的
# 所以需要转换为 bacth * sequence
return Batch(src, trg, pad_idx)
最后终于到了 Batch 类, 最关键的信息来⾃这⾥:
class Batch:
"Object for holding a batch of data with mask during training."
def __init__(self, src, trg=None, pad=0):
self.src = src
self.src_mask = (src != pad).unsqueeze(-2)
# 在预测的时候是没有 tgt 的，此时为 None
if trg is not None:
< = trg[:, :-1]
# 每次迭代的时候, 去掉最后⼀个单词
<_y = trg[:, 1:]
# 去掉第⼀个单词
<_mask = self.make_std_, pad)
# target 语⾔中单词的个数
@staticmethod
def make_std_mask(tgt, pad):
"Create a mask to hide padding and future words."
tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & transformer.subsequent_mask(tgt.size(-1)).type_as(tgt_mask)
# tgt.size(-1) 表⽰的是序列的长度
return tgt_mask
在 class Batch 中, trg 为 None 的时候很好理解, 也就是在预测的时候, 是没有⽬标语⾔的, 其实在预测
的时候, 只有输⼊的 Batch , 那么预测过程的 Attention Mask ⼜是如何实现的呢? 这个我们放在后⾯再说, 先看这⾥的src_mask, 源语⾔的 mask, 也就是 encoder 时的self_Attention 时的mask, 这个很好理解, 就是将⾮ 0 的数字变成 1, 获得⼀个 0/1 矩阵, = trg[:, :-1]这⾥去掉的最后⼀个单词, 不是真正的单词, ⽽是标志'<eos>' , 输⼊与输出都还有⼀个 '<sos>' 在句⼦的开头, _y = trg[:, 1:]去掉开头就变成了最后的结果. 接下来就是最关键的获取 target 语⾔的 mask 矩阵,
def subsequent_mask(size):
decoder
"Mask out subsequent positions."
attn_shape = (1, size, size)
subsequent_mask = np.s(attn_shape), k=1).astype('uint8')
return torch.from_numpy(subsequent_mask) == 0
这个函数⼲了啥呢?
我们先写成这样:
def subsequentmask(size):
"Mask out subsequent positions."
attn_shape = (1, size, size)
subsequent_mask = np.s(attn_shape), k=1).astype('uint8')
return subsequent_mask == 0
print(subsequentmask(5))
>>
[[[ True False False False False]
[ True  True False False False]
[ True  True  True False False]
[ True  True  True  True False]
[ True  True  True  True  True]]]
当这个 numpy 数组转化为tensor 的时候, 构成的是维度为 (1, 5, 5) 的矩阵. 我们注意到, self.src_mask = (src != pad).unsqueeze(-2)也就是说, 源语⾔的 mask 矩阵的维度是 (batch_size, 1, length), 那么为什么attn_shape = (batch_size, size, size)呢? 可以这么解释, 在 encoder 阶段的
Self_Attention 阶段, 所有的 Attention 是可以同时进⾏的, 把所有的 Attention_result 算出来, 然后⽤同⼀个 mask vector * Attention_result 就可以了, 但是在 decoder 阶段却不能这么做, 我们需要关注的问题是:
根据已经预测出来的单词预测下⾯的单词, 这⼀过程是序列的,
⽽我们的计算是并⾏的, 所以这⼀过程中, 必须要引⼊矩阵. 也就是上⾯的 subsequent_mask() 函数获得的矩阵.
这个矩阵也很形象, 分别表⽰已经预测的单词的个数为, 1, 2, 3, 4, 5.
然后我们将以上过程反过来过⼀遍, 就很明显了, 在 batch阶段获得 mask 矩阵, 然后和 batch ⼀起训练, 在 encoder 与 deocder 阶段实现mask 机制.
预测过程中的 Mask实现
我们直接来看预测过程中的 decoder 的实现,
def greedy_decode(model, src, src_mask, max_len, start_symbol):
memory = de(src, src_mask)
# memory 是 encoder 的中间结果
batch_size = src.shape[0]
ys = s(batch_size, 1).fill_(start_symbol).type_as(src)
# 预测句⼦的初始化
for i in range(max_len-1):
out = model.decode(memory, src_mask, ys, transformer.subsequent_mask(ys.size(1)).type_as(src))
# ys 的维度是 batch_size * times, 所以target_mask 矩阵必须是 times * times
# 根据 decoder 的训练步骤, 这⾥的 out 输出就应该是 batch_size * (times+1) 的矩阵
prob = ator(out[:, -1])
# out[:, -1] 这⾥是最新的⼀个单词的 embedding 向量
# generator 就是产⽣最后的 vocabulary 的概率, 是⼀个全连接层
_, next_word = torch.max(prob, dim = 1)
# 返回每⼀⾏的最⼤值, 并且会返回索引
next_word = next_word.unsqueeze(1)
ys = torch.cat([ys, pe_as(src)], dim=1)
# 将句⼦拼接起来
return ys
上⾯代码的transformer.subsequent_mask(ys.size(1)).type_as(src)这⼀部分就很好的解释了 target_mask 矩阵的构造⽅法, 在这⾥, 输⼊不是像训练时候⼀样, 是整个的⽬标语⾔的句⼦, ⽽是已经训练的句⼦的集合, 这⾥的 decode 的步骤不仅仅是预测了最后⼀个单词, 同时, 前⾯所有的单词都进⾏了预测, 只是预测的结果和上次是⼀样的. 这⾥在多说⼀点, 在预测的时候,

版权声明：本站内容均来自互联网，仅供演示用，请勿用于商业和其他非法用途。如果侵犯了您的权益请与我们联系QQ:729038198，我们将在24小时内删除。