訓練過程中的 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, key.transpose(-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.encoder = encoder
self.decoder = decoder
self.src_embed = src_embed
# 將源語言的單詞 embedding 放在一起, position embedding
self.tgt_embed = tgt_embed
# 將目標語言的單詞 embedding 放在一起, position embedding
self.generator = generator
# 就是最后產生結果的地方
def forward(self, src, tgt, src_mask, tgt_mask):
"Take in and process masked src and target sequences."
return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
def encode(self, src, src_mask):
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_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):
model.train()
out = model.forward(batch.src, batch.trg ,batch.src_mask, batch.trg_mask)
# 參數來自 batch
loss = loss_compute(out, batch.trg_y, batch.ntokens)
# 這一步既計算了損失, 又更新了參數
total_loss += loss
total_tokens += batch.ntokens
tokens += batch.ntokens
這些都是訓練的步驟, 數據是怎么來的, mask 矩陣來自 batch, 所以最關鍵的是 batch 是怎么來的, 再往回找在 train.py函數中, 我們發現
_, logger_file = train_utils.run_epoch(args, (train_utils.rebatch(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 = batch.src.transpose(0, 1), batch.trg.transpose(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:
self.trg = trg[:, :-1]
# 每次迭代的時候, 去掉最后一個單詞
self.trg_y = trg[:, 1:]
# 去掉第一個單詞
self.trg_mask = self.make_std_mask(self.trg, pad)
self.ntokens = (self.trg_y != pad).sum().item()
# 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 矩陣, self.trg = trg[:, :-1] 這里去掉的最后一個單詞, 不是真正的單詞, 而是標志 '<eos>' , 輸入與輸出都還有一個 '<sos>' 在句子的開頭, self.trg_y = trg[:, 1:] 去掉開頭就變成了最后的結果. 接下來就是最關鍵的獲取 target 語言的 mask 矩陣,
def subsequent_mask(size):
"Mask out subsequent positions."
attn_shape = (1, size, size)
subsequent_mask = np.triu(np.ones(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.triu(np.ones(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 = model.encode(src, src_mask)
# memory 是 encoder 的中間結果
batch_size = src.shape[0]
ys = torch.ones(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 = model.generator(out[:, -1])
# out[:, -1] 這里是最新的一個單詞的 embedding 向量
# generator 就是產生最后的 vocabulary 的概率, 是一個全連接層
_, next_word = torch.max(prob, dim = 1)
# 返回每一行的最大值, 並且會返回索引
next_word = next_word.unsqueeze(1)
ys = torch.cat([ys, next_word.type_as(src)], dim=1)
# 將句子拼接起來
return ys
上面代碼的 transformer.subsequent_mask(ys.size(1)).type_as(src) 這一部分就很好的解釋了 target_mask 矩陣的構造方法, 在這里, 輸入不是像訓練時候一樣, 是整個的目標語言的句子, 而是已經訓練的句子的集合, 這里的 decode 的步驟不僅僅是預測了最后一個單詞, 同時, 前面所有的單詞都進行了預測, 只是預測的結果和上次是一樣的. 這里在多說一點, 在預測的時候,