Attention is all you need 源碼解析
最近學習Transformer模型的時候，並且好好讀了一下Google的《Attention is all you need》論文。論文地址如下: Attention is All you need. 同時學習了一下其github的代碼，代碼地址如下：github code. 在網上查資料的過程中，還找到了一個好像也用的比較多的版本：Transformer demo.
在網上學習的過程中，我發現有一些csdn blog的源碼解析內容主要是根據後一個的內容來進行的。在對比兩個github code 的過程中，我感覺前一個的內容質量可能更好，代碼也更完整一些而且不斷有人在做更新，所以本次解析也主要以第一個github code 的版本來進行，同時也是希望在剛調通完代碼後可以自己記錄一些在學習代碼過程中遇到的一些理解方面的問題，進行記錄，也方便後續的回顧。
模型中某些理論方面的內容，可以參考我之前寫的Transformer模型中Encoder和Decoder各結構的理解
prepro.py
該模塊用於一開始的加載原始數據 -> 對數據進行處理 -> 對句子進行切分。
_prepro利用lambda的形式設計函數，第一個_prepro用於去除train.tags.de-en.de和train.tags.de-en.en文件中以<開始的sentence. 第二個_prepro利用了Python的re正則表達式模塊，<[^>]+>，其中[^>]+表示包含除了右括號以外的其他字符。表示的意思是：選擇以左括號開始，右括號結束的部分。
之後即是利用_write()將對應進行處理後的數據集寫入到*iwslt2016/prepro/*目錄下。
這裏在講一下sentencepiece模塊。sentencepiece是一個Google開源的自然語言處理包，可以對分詞結果進行訓練的一個模塊。我們知道，在不同的領域或不同的句子下，我們希望某些詞的分詞效果有所不同(包括該領域的特殊詞彙). 例如，“抗精神病治療”，我們希望在分詞切分的時候可以以整體的形式出現，而不是被切分成“抗”，“精神病”，“治療”。在現有的某些分詞工具下，例如jieba分詞，對於以上情況我們可以添加自定義詞典的形式來加入一直詞。但有個問題就是我們無法一開始獲知所有的組合，即某些詞彙組合我們可能是未知的，所以我們就希望機器能夠自動學習經常出現的短語或詞，Sentencepiece就是來解決這個問題的。我們可以用該領域內的大量文本進行訓練，該模塊可以根據樣本訓練的結果自動生成詞表和分詞模型。然後我們可以根據模型來對文本進行分詞的工作。
有興趣大家可以自己拿一個小文本作爲一個小demo來嘗試一些訓練、加載模型和分詞的過程。
更詳細的內容可以參考以下鏈接內容：
自然語言處理之_SentencePiece分詞
Unsupervised text tokenizer for Neural Network-based text generation.
SentencePiece Python Wrapper
loaddata.py
該模塊用於構建數據集、對數據集根據詞表進行轉換等功能。
load_vocab()
函數即根據已構建的詞表來構建idx2token和token2idx兩個映射字典。
load_data()
函數即對輸入文件(source file)和目標文件(target file)的sentence進行處理，只保留每行sentence的長度小於等於maxlen的句子。並以列表的形式進行保存每個滿足條件的sentence.
encode()
函數即根據idx2token和token2idx映射表，將每個單詞轉換成對應的idx表示。同時，這裏對encode部分輸入的sentence末尾加入了</s>結束符，在decode部分的sentence開頭加入<s>表示解碼開始，</s>表示解碼結束。
generate_fn()
函數即以Python生成器的形式，生成training / evaluation data. yield (x, x_seqlen, sent1), (decoder_input, y, y_seqlen, sent2)包含編碼器和解碼器的數據兩部分。第一部分包含三個元素：x代表由Idx表示的sentence，維度大小爲(N,T1),N表示batch_size, T1表示句子的長度。x_seqlen表示句子的長度，維度爲(N，)，sent1表示原始由token表示的sentence，維度爲（N，）。第二部分包含四個元素，其內容與編碼器中的部分類似，只是decoder_inputs表示不包括結束符</s>，y表示不包括開始符<s>。
input_fn()
函數中利用tf.data.Dataset.from_generator()函數來返回一個數據的生成器對象。自Tensorflow 1.x以來，逐步開發引入了tf.data.Dataset模塊，使其數據讀入的操作變得更爲方便。tf.data.Data模塊讀入數據也有兩種方式，一種是tf.data.Data.from_tensor_slices()另一種是tf.data.Data.from_generate().兩種方法都是對x,y兩部分數據沿第一個維度進行切分然後進行組合成tuple。
from_tensor_slices()是直接利用已存在的兩組數據進行切分組合,每次分別從inputs和labels各取batch_size個大小的數據進行組合。在不加repeat()的情況下爲不重複抽取，即只能迭代原數據集中batch大小 / batch_size大小的次數。如果想迭代任意次數，則需要加入.repeat()方法
from_generator()是按生成器對象，按batch_size大小逐batch讀入數據。對每次生成的inputs,labels進行組合成一個tuple
from_tensor_slices()是當數據集已經生成好的時候比較適用，當數據集是需要動態生成過程中時，則可以使用from_generator()函數，每次動態的從生成器中生成指定大小的數據集.
詳細內容可以參考：
使用TensorFlow Dataset讀取數據
get_batch()
函數即獲取training / evaluation mini-batch。利用input_fn()函數返回數據集生成器對象batches，並根據樣本總數和batch_size的大小計算出所需要的batch數目num_batches，以及樣本總數len(sents1)。
modules.py
該模塊用於構建Transformer模型中的各部分的模型結構，即論文中的 3 Model Architecture部分，根據論文中模型的結構和參數來構建相應的模型結構。這一部分的代碼即實現Transformer中的模型結構，所以會重點講解這部分的代碼內容。
layer_normalizaiton()
函數即是實現模型中的Layer Normalizaiton. 由layer normalizaiton的原理可知，我們需要計算的是在某一層上的mean和variance，比如在mulit-head attention之後的LN, LN函數的輸入Inputs的維度爲(N, T_q, d_model)，則此時需要在d_model的維度上求mean和variance。所以這裏mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)中 axis=-1。接下來的beta、gamma兩個參數則是normalization的線性調整參數，可以讓分佈的曲線壓縮或延長一點，左移或右移一點的進行調整：outputs = gamma * normalized + beta。最後返回經過標準化後的數據。
get_token_embeddings()
函數是初始化word embedding的部分。這裏對embedding的數值利用了xavier_initializer()進行初始化。同時要注意一個zero_pad的參數，該參數是爲了使queries / keys的mask更加方便。
scaled_dot_product_attention()
函數即模型的single self-attention結構。按照
的公式進行計算。其中causality參數即表示是否applies masking for future blinding. 該算法的每一步即是按照論文中的scaled Dot-Product Attention進行的
mask()
這裏要重點分析一下mask的函數內容。函數中設置了三個判斷條件，分別是要在keys有padding、queries有padding和是否mask for future blinding。這裏對Keys和queries分情況討論的原因是，因爲在不是self-attention中，query和key的第二維可能會有不同（兩者的句子長度可能會不同），即可能一個有pad一個沒有pad，所以要分情況討論。對於keys和queries兩種情況，以keys的情況舉例來說明：
這裏tf.reduce_sum(tf.abs(keys))的目的是爲了在keys的tensor中，查找出那些padding項。padding項由於一開始做pad時，其embedding設置爲全0，所以對其embedding(d-model)維度進行求和，若是padding項，則其和只會爲0.然後再又tf.sign進行二分，即最終結果只會出現0，1兩項。padding項其求和恆爲0，非padding項其求和>0則爲1.
此時，mask的維度由(N,T_k,d)降爲(N,T_k).其中，T_k的維度中的0，1代表了哪些位置是padding哪些位置是word。由於我們後續在利用inputs對paddings進行替換的時候，tf.where(input,a,b)要求tensor a,b要有一樣的尺寸，所以需要對maks的維度進行擴充。
這裏對於爲什麼擴充重複的維度爲tf.shape(queries)[1]可以這樣理解爲：
首先，inputs的維度爲(N,T_q, T_k),代表的即是query對每個key下的attention score的分數矩陣。T_q代表每一個query，T_k代表每一個key。由以上mask的分析可知，mask得出的即是在T_k下有哪些是padding，哪些是word，那我們就要對T_k下爲0所代表的列下的attention score進行替換。此時，通過tf.expand_dims(masks,1),mask的維度變爲了(N,1,T_k), 相當於是只有在inputs分數矩陣中的第一行標明瞭哪些是word哪些是padding。（可以理解爲mask維度中第二維代表的是在inputs矩陣中的行表示query，T_k代表的列表示key）由於要對inputs中的每一行中，key爲Padding的位置都進行替換，所以第二維度要擴展query的個數（行數）。
接下來就是padding設置爲和inputs一樣的大小（因爲要利用Inputs對padding進行替換）。其值設置爲非常小，目的是爲了減少padding位置的attention score值對後續value進行的影響。tf.equal（）來判斷mask中哪些位置爲0，爲0的變爲1.（這裏爲1代表了padding的位置）。tf.where()表示以paddings爲基準，對與mask中爲False的位置(映射到paddings位置上)，將inputs的對應位置的元素替換padding中對應位置的元素，得到的結果即爲，保留了原始paddings中的極小值部分(這部分即爲padding項)，其他位置替換爲inputs中的元素。
最後生成的維度和inputs一樣爲（N,T_q,T_k）的attention score矩陣。
masks = tf.sign(tf.reduce_sum(tf.abs(keys), axis = -1)) #(N, T_k)
masks = tf.expand_dims(masks, 1) #(N, 1, T_k)
masks = tf.tile(masks, [1, tf.shape(queries)[1], 1]) #(N, T_q, T_k)
#Apply masks to inputs
paddings = tf.ones_like(inputs) * padding_num
outputs = tf.where(tf.equal(masks, 0), paddings, inputs) #(N, T_Q, T_k)
queries的部分和keys的相同。接下來再分析一下mask for future blinding時的情況：
這裏是對decoder中的Mask Multi-head attention. type表示是否要屏蔽未來的信息。
由於在解碼器的部分，不像編碼器一樣可以直接並行運行，解碼器由於需要翻譯，仍然是序列化的進行。所以在decoder的Attention部分(這裏其實也是self-Attention)，某一時刻翻譯出來的詞只能跟它之前翻譯出的詞產生聯繫。即只能跟它之前的詞計算Attention score.所以這裏就用了一個下三角矩陣。
這樣在計算Attention時，前面的word不會與在它後面的word產生聯繫。比如第一個word在計算Matmul(V)時，就只與它自己相關。（即不會有跟它後面的詞的attention score）
diag_vals = tf.ones_like(inputs[0, :, :]) #(T_q, T_k)
tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense() #(T_q, T_k).生成一個下三角矩陣
masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(inputs)[0],1,1]) #(N, T_q, T_k)
paddings = tf.ones_like(masks) * padding_num
outputs = tf.where(tf.equal(masks, 0), paddings, inputs)
multihead_attention()
函數即是多頭注意力機制的函數實現。
#Linear projections
Q = tf.layers.dense(queries, d_model, use_bias = False) #(N, T_q, d_model)
K = tf.layers.dense(keys, d_model, use_bias = False) #(N, T_k, d_model)
V = tf.layers.dense(values, d_model, use_bias = False) #(N, T_k, d_model)
即先經過一個線性映射將queries、keys、values從d_k，d_v的維度映射到d_model的維度，接下來這裏的多頭注意力機制並不是分開的8個tensor，而是對原始的Q、K、V矩陣進行切分在進行拼接而成的。具體的過程如下：
首先對於d_model = 512，由於採用了num_heads = 8，則正如論文中所說的，在每一個head下，d_k = d_v = d_model / 8 = 64。所以一開始tf.split()函數的axis = 2，即沿d_model維度進行切分，切分成8片。然後對每片，沿第一個維度batch_size的維度進行拼接，即形成了維度爲（h*N, T_q, d_model/h）的維度（針對矩陣Q而言）。該維度可以明顯的看出，生成了h個大小爲(N,，T_q，d_model/h)的矩陣Q_。對於矩陣K，V的分析則類似。
#這裏axis所在維度的長度 / num_heads 應該能被整除。
Q_ = tf.concat(tf.split(Q, num_heads, axis = 2), axis = 0) #(h*N, T_q, d_model/h)
K_ = tf.concat(tf.split(K, num_heads, axis = 2), axis = 0) #(h*N, T_k, d_model/h)
V_ = tf.concat(tf.split(V, num_heads, axis = 2), axis = 0) #(h*N, T_k, d_model/h)
在對這h個Q_、K_、V_矩陣做scaled_dot_product_attention之後，再進行Reshape的操作，即做跟以上切分相反的操作：先按axis = 0第一個維度做切分，相當於生成h個維度大小爲(N，T_q，d_model/h)的矩陣，然後再對這h個矩陣按axis = 2即按（d_model/h）的維度進行拼接，從而重新生成大小爲（N，T_q，d_model）的矩陣。
outputs = tf.concat(tf.split(outputs, num_heads, axis = 0), axis = 2) #(N, T_q, d_model)
feed_forward()
函數即是實現全連接的前向神經網絡的部分。代碼中利用了tf.layers.dense()函數來實現全連接的網絡，第一個代表的從輸入層到第一隱層的過程，所以num_hiddens設置爲num_units[0]（論文中d_ff = 2048），第二個則是由第一隱層到輸出層，所以num_hiddens設置爲num_units[1]（論文中爲d_model = 512）。
Inner layer
outputs = tf.layers.dense(inputs, num_units[0], activation = tf.nn.relu)
#Outer layer
outputs = tf.layers.dense(outputs, num_units[1])
label_smoothing()
這部分就相當於是使矩陣中的數進行平滑處理。把0改成一個很小的數，把1改成一個比較接近於1的數。論文中說這雖然會使模型的學習更加不確定性，但是提高了準確率和BLEU score。（論文中的5.4部分）
position_embedding()
#Second part, apply the cosine to even columns and sin to odds
#list[start : end : step]
position_enc[:, 0::2] = np.sin(position_enc[:,0::2]) #dim 2i.取下標爲偶數的
position_enc[:, 1::2] = np.cos(position_enc[:,1::2]) #dim 2i+1.取下標爲奇數的
position_enc = tf.convert_to_tensor(position_enc, tf.float32) #(maxlen, E)
這裏即按照論文中的，對於下標爲偶數的則採用sin，對下標爲奇數的則採用cos的方式。position_enc[:，0::2]中的0::2代表的意義即是從第一個元素開始(0)每隔2個位置進行遍歷，即取了下標爲偶數的位置。
#從position_enc 中 取position_ind對應索引位置的元素
outputs = tf.nn.embedding_lookup(position_enc, position_ind)
model.py
該部分模塊即利用上文已經構建的各組件，來構建整體的Transformer模型結構。
這裏代碼中也標明瞭xs，ys所代表的意義。也是看上文中load_data.py文件中的generate_fn()函數的數據生成器部分。
__init__
首先__init__函數爲Transformer包含了一個hyperparams的超參數對象，來生成token2idx和idx2token的映射，以及word embedding 矩陣。
encode()
encode()函數即Transformer模型的編碼器部分，按照論文中的模型結構來進行搭建。首先是word embedding + position embedding來形成輸入數據。接下來是一個循環num_blocks次的Multi-Head Attention部分。利用modules.py模塊中的multihead_attention()函數來進行搭建，其中causality參數=False代表只是mask和不是mask for future blinding。multi-head attention之後再接一個Feed Forward的前向網絡，由這兩個sub_layers構成一個block。
decode()
decode()函數即Transformer模型的解碼器部分，按照論文中的模型結構來進行搭建。首先仍然是一個embedding來構建輸入數據的部分，接下來是一個循環num_blocks次的部分。blocks中的第一個是一個Mask Multi-Head Attention，所以需要在multihead_attention()函數中將causality設置爲True。又因爲該Attention是self-Attention，所以輸入的queries、keys、values都是dec本身。
接下來仍然是一個Multi-Head Attention,這裏的keys、values由encoder的輸出提供，所以輸入的參數與上面的不同。這裏causality參數設置爲False。然後是一個同樣的Feed Forward的前向網絡
最後是一個decoder解碼器的輸出部分。這裏linear projection的權重矩陣由embedding矩陣的轉置得到，因爲最終輸出要生成的是一個vocab_size大小的向量，表示輸出的各個單詞的概率。這裏Multi-Head Attention的輸出與權重矩陣的乘法，沒有使用tf.matmul()的矩陣乘法，因爲這裏兩個矩陣的維度大小不相同。Multi-Head Attention的輸出維度大小爲（N，T2，d_model），而weights的維度大小爲（d_model，vocab_size）。所以這裏用了tf.einsum()函數，這裏該函數的第一個參數：'ntd,dk->ntk’表示的意思是，->代表乘法操作，ntd,dk->ntk表示兩個矩陣相乘後結果的維度爲ntk。這樣就實現了兩個維度不同矩陣的乘法。
liner projuction的輸出logits的維度大小爲（N，T2，vocab_size）。即代表了在當前T2的長度下，每個位置上的輸出各個單詞的概率。然後利用tf.argmax() 在axis = -1的維度下，求出概率最大的那個位置的詞作爲該位置上的輸出單詞。所以y_hat的維度大小爲（N，T2）。
train()
該函數即是進行模型訓練的部分。首先是調用encode()和decode()函數來獲取各部分的輸出結果。接下來就是模型訓練的scheme.
利用one_hot表示每個詞的索引在整個詞表中的位置,相當於構建出了要訓練的目標Label，這裏就是要使logits的最終結果，即vocab_size大小的向量中，目標詞彙所在位置（索引）的值儘可能的大，而使其他位置的值儘可能的小。構造出了輸出和標籤之後，就使用tf.nn.softmax_cross_entropy_with_logits()進行訓練。
y_ = label_smoothing(tf.one_hot(y, depth = self.hp.vocab_size)) #(N, T2, vocab_size)
ce = tf.nn.softmax_cross_entropy_with_logits(logits = logits, labels = y_) #(N,T2)
在計算Loss之前，還要進行一定的處理。由於一開始對有些句子長度不夠maxlen的進行了padding，所以在計算Loss的時候，將這些位置上的誤差清0
nonpadding = tf.to_float(tf.not_equal(y, self.token2idx['<pad>'])) #0:<pad>
loss = tf.reduce_sum(ce * nonpadding) / (tf.reduce_sum(nonpadding) + 1e-7) #最終求得一個loss scalar進行優化.
最後即是利用了AdadeltaOptimizer優化器對loss進行優化。tf.summary.scalar()函數即是以key-value的形式保存數值，用於TensorBoard中對數據的可視化展示。
eval()
該函數是對模型訓練效果進行評估的模塊。這裏即是讓解碼器根據開始符<s>來自動的進行Machine Translation的過程。流程與test的部分類似，但這裏並不是真正的test。
decoder_inputs = tf.ones((tf.shape(xs[0])[0],1), tf.int32) * self.token2idx['<s>']
ys = (decoder_inputs, y, y_seqlen, sents2)
這裏xs[0]表示取第一個tensor x。tf.shape(xs[0]) = [N,T], tf.shape([xs[0])[0] 即取batch_size大小，在evaluation部分，由於解碼器的預測仍然是按序列式的進行（與train時候的不同）,即是每一次解碼過程預測一個目標詞彙，所以在時刻t=0時解碼器的輸入維度應該是(N,1)，即此時爲一個batch輸入，每個batch的開頭爲<s>表示開始進行解碼，然後每完成一次解碼過程，則每個batch已輸出詞彙數+1，例如t=1時刻，則解碼器的輸入維度爲(N,2),以此類推，直到輸入到表示停止。然後再將新的decoder_inputs加入ys中作爲下一時刻decoder的輸入。
for _ in tqdm(range(self.hp.maxlen2)):
logits, y_hat, y, sents2 = self.decode(ys, memory, False)
if tf.reduce_sum(y_hat, 1) == self.token2idx['<pad>']:break
_decoder_inputs = tf.concat((decoder_inputs, y_hat), 1)
ys = (_decoder_inputs, y, y_seqlen, sents2)
這裏這一過程，即是不斷的進行序列化的預測過程。循環次數爲maxlen2次，表示是要翻譯完一整個句子的長度，然後不斷的將上一時刻的解碼器的輸出添加到下一時刻解碼器的輸入。
#monitor a random sample
n = tf.random_uniform((), 0, tf.shape(y_hat)[0]-1, tf.int32) #即從0，batch_size-1之間選擇一個batch sample進行觀察
sent1 = sents1[n] #input sentence
pred = convert_idx_to_token_tensor(y_hat[n], self.idx2token) #prediction sentence
sent2 = sents2[n] #output sentence
這裏即是隨機抽取一個batch來查看模型的結果。n代表從0，batch_size-1之間選擇一個batch sample進行觀察。sent1[n]表示原始的輸入句子，pred即代表了decoder的預測翻譯的輸出句子，sents2[n]即表示正確的翻譯輸出句子。
train.py
該模塊即是設計對模型進行訓練的主方法。首先，即是調用get_batch()函數來生成訓練和評估時候的數據。
#create a iterator of the correct shape and type
#A reinitializable iterator is defined by its structure.We could use the 'output_types' and 'output_shapes' properties of train_batches.
iter = tf.data.Iterator.from_structure(train_batches.output_types, train_batches.output_shapes)
xs, ys = iter.get_next()
#即利用已有的dataset對象，來初始化一個新的數據集生成器
train_init_op = iter.make_initializer(train_batches)
eval_init_op = iter.make_initializer(eval_batches)
這裏iter利用了from_structure()方法來實現，該函數的參數即是生成數據的類型和大小，這裏根據了train_batches的類型和大小。然後分別利用train_batches和eval_batches來初始化訓練和評估的數據集生成器。
m = Transformer(hp)
loss, train_op, global_step, train_summaries = m.train(xs,ys)
y_hat, eval_summaries = m.eval(xs,ys)
這裏即是加載模型，然後調用模型裏的train()和eval()方法來進行訓練和做評估。
接下來就是構建Tensorflow的Session來對整個模型進行訓練和評估的過程：
ckpt = tf.train.latest_checkpoint(hp.logdir)用來查找到最近的檢查點文件。
save_variable_specs(os.path.join(hp.logdir, "specs"))則是用來保存訓練過程中的一些參數變量。
summary_writer = tf.summary.FileWriter(hp.logdir, sess.graph)即是要利用TensorBoard來進行數據可視化展示。
sess.run(train_init_op)來運行一次數據集生成器，即生成一次數據集。
total_steps = hp.num_epochs * num_train_batches這裏即循環epochs次，每epoch次要對num_train_batches個batch進行訓練，也就是每次epoch都要對所有的batch進行一次訓練，以此來計算總的計算次數。
_, _gs, _summary = sess.run([train_op, global_step, train_summaries])這一步即是當一個epoch沒有計算完時（即所有的batch還沒有計算完時），不斷的迭代去訓練train_op。每運行一次該行代碼，都會調用一次xs, ys = iter.get_next()來生成一個batch數據集。
if _gs and _gs % num_train_batches == 0:這裏即是每次epoch全部訓練完所有batch，則打印相應的信息，並對模型效果進行測試。
if語句以下的內容包括在評估中將預測輸出由idx轉化爲對應的token，然後將結果寫入到本地，計算bleu的得分，保存模型的檢查點文件。
總結
比較重要的幾個模塊文件差不多就是上面的這幾個模塊。學習Transformer源代碼的過程中，也讓我對Tensorflow構建模型的過有了新的認識。如果有寫的不對的地方還希望大家指正，希望能夠與大家一起交流學習。