一文讀懂 BERT 源代碼
大數(shù)據(jù)文摘授權(quán)轉(zhuǎn)載自數(shù)據(jù)派THU
作者：陳之炎
BERT模型架構(gòu)是一種基于多層雙向變換器（Transformers）的編碼器架構(gòu)，在tensor2tensor庫框架下發(fā)布。由于在實(shí)現(xiàn)過程當(dāng)中采用了Transformers，BERT模型的實(shí)現(xiàn)幾乎與Transformers一樣。
BERT預(yù)訓(xùn)練模型沒有采用傳統(tǒng)的從左到右或從右到左的單向語言模型進(jìn)行預(yù)訓(xùn)練，而是采用從左到右和從右到左的雙向語言模型進(jìn)行預(yù)訓(xùn)練，本文對BERT模型預(yù)訓(xùn)練任務(wù)的源代碼進(jìn)行了詳細(xì)解讀，在Eclipse開發(fā)環(huán)境里，對BERT 源代碼的各實(shí)現(xiàn)步驟分步解析。
BERT 模型的代碼量比較大，由于篇幅限制，不可能對每一行代碼展開解釋，在這里，解釋一下其中每一個核心模塊的功能。
1) 數(shù)據(jù)讀取模塊
圖 1
模型訓(xùn)練的第一步，是讀取數(shù)據(jù)，將數(shù)據(jù)從數(shù)據(jù)集中讀取進(jìn)來，然后按照BERT 模型要求的數(shù)據(jù)格式，對數(shù)據(jù)進(jìn)行處理，寫出具體數(shù)據(jù)處理的類以及實(shí)際要用到的數(shù)據(jù)集中數(shù)據(jù)處理的方法，如果任務(wù)中用到的數(shù)據(jù)集不是MRPC ，這部分的代碼需要依據(jù)特定的任務(wù)重新寫一下如何操作數(shù)據(jù)集的代碼，對于不同的任務(wù)，需要構(gòu)造一個新的讀取數(shù)據(jù)的類，把數(shù)據(jù)一行一行地讀進(jìn)來。
2) 數(shù)據(jù)預(yù)處理模塊
圖 2
利用tensorflow 對數(shù)據(jù)進(jìn)行預(yù)處理，由于用TF-Record 讀數(shù)據(jù)的速度比較快，使用起來比較方便，在數(shù)據(jù)讀取層面，需要將數(shù)據(jù)轉(zhuǎn)換成TF-Record格式。首先，定義一個writer,利用writer函數(shù)將數(shù)據(jù)樣本寫入到TF-Record當(dāng)中，這樣一來，在實(shí)際訓(xùn)練過程中，不用每次都到原始數(shù)據(jù)中去讀取數(shù)據(jù)，直接到TF-Record當(dāng)中讀取處理好的數(shù)據(jù)。
把每一個數(shù)據(jù)樣本都轉(zhuǎn)化成一個TF-Record格式的具體做法如下：首先，構(gòu)建一個標(biāo)簽，接下來對數(shù)據(jù)做一個判斷，判斷數(shù)據(jù)中由幾句話組成，拿到當(dāng)前第一句話后，做一個分詞操作。分詞方法為wordpiece 方法。在英文文本中，由字母組成單詞，詞與詞之間利用空格來切分單詞，利用空格切分單詞往往還不充分，需要對單詞做進(jìn)一步切分轉(zhuǎn)換，在BERT 模型中，通過調(diào)用wordpiece 方法將輸入的單詞再進(jìn)一步切分，利用wordpiece的貪心匹配方法，將輸入單詞進(jìn)一步切分成詞片，從而使得單詞表達(dá)的含義更加豐富。在這里，利用wordpiece 方法將讀入的單詞進(jìn)行再次切分，把輸入的單詞序列切分成更為基本的單元，從而更加便于模型學(xué)習(xí)。
在中文系統(tǒng)中，通常把句子切分成單個的字，切分完成之后，把輸入用wordpiece轉(zhuǎn)化成wordpiece結(jié)構(gòu)之后，再做一個判斷，看是否有第二句話輸入，如果有第二句話輸入，則用wordpiece對第二句話做相同的處理。做完wordpiece轉(zhuǎn)換之后，再做一個判斷，判斷實(shí)際句子的長度是否超過max_seq_length 的值，如果輸入句子的長度超過max_seq_length規(guī)定的數(shù)值，則需要進(jìn)行截斷操作。
3) tf-record 制作
對輸入句對進(jìn)行編碼，遍歷wordpiece結(jié)構(gòu)的每一個單詞以及每一個單詞的type_id ,加入句子分隔符【CLS】、【SEP】，為所有結(jié)果添加編碼信息；添加type_id，把所有單詞映射成索引功，對輸入詞的ID （標(biāo)識符）進(jìn)行編碼，以方便后續(xù)做詞嵌入時候進(jìn)行查找；
Mask編碼：對于句子長度小于max_seq_length 的句子做一個補(bǔ)齊操作。在self_attention 的計算中，只考慮句子中實(shí)際有的單詞，對輸入序列做input_mask 操作，對于不足128個單詞的位置加入額外的mask,目的是讓self_attention知道，只對所有實(shí)際的單詞做計算，在后續(xù)self_attention計算中，忽略input_mask=0 的單詞，只有input_mask=1 的單詞會實(shí)際參與到self_attention計算中。Mask編碼為后續(xù)的微調(diào)操作做了初始化，實(shí)現(xiàn)了任務(wù)數(shù)據(jù)的預(yù)處理。
圖 3
對input_Feature做初始化：構(gòu)建 input_Feature并把結(jié)果返回給BERT。通過一個for 循環(huán)，遍歷每一個樣本，再對構(gòu)造出來一些處理，把input_id、input_mask和segment_id均轉(zhuǎn)換成為int類型，方便后續(xù)tf-record的制作。之所以要做數(shù)據(jù)類型的轉(zhuǎn)換，是因?yàn)閠ensorflow 官方API要求這么做，tensorflow對tf-record的格式做了硬性的規(guī)定，用戶無法自行對其修改。在后續(xù)具體項(xiàng)目任務(wù)中，在做tf-record時，只要把原始代碼全部復(fù)制過去，按照原有的格式修改即可。構(gòu)造好input_Feature之后，把它傳遞給tf_example,轉(zhuǎn)換成tf_train_features ，之后，直接寫入構(gòu)建好的數(shù)據(jù)即可。
4) Embeding層的作用
在BERT 模型中有一個creat_model 函數(shù)，在creat_model 函數(shù)中一步一步把模型構(gòu)建出來。首先，創(chuàng)建一個BERT 模型，該模型中包含了transformer的所有結(jié)構(gòu)，具體操作過程如下：
圖 4
讀入配置文件，判斷是否需要進(jìn)行訓(xùn)練，讀入input_id、input_mask和segment_id等變量, one_hot_embedding變量?在利用TPU 訓(xùn)練時才使用，在用CPU 訓(xùn)練時不用考慮，默認(rèn)值設(shè)為Faulse。
構(gòu)建embedding層，即詞嵌入，詞嵌入操作將當(dāng)前序列轉(zhuǎn)化為向量。BERT 的embedding層不光要考慮輸入的單詞序列，還需要考慮其它的額外信息和位置信息。BERT 構(gòu)建出來的詞嵌入向量中包含以下三種信息：即輸入單詞序列信息、其它的額外信息和位置信息。為了實(shí)現(xiàn)向量間的計算，必須保持包含這三種信息的詞向量的維數(shù)一致。
5) 加入額外編碼特征
接下來，進(jìn)入到embedding_lookup 層，這個層的輸入是：input_id（輸入標(biāo)識符）、vocab_size(詞匯表大小)、embedding_size(詞嵌入的維度)、initializer_range(初始化的取值范圍)。embedding_lookup的輸出是一個實(shí)際的向量編碼。
圖 5
首先，獲取embedding_table，然后到embedding_table里查找每個單詞對應(yīng)的詞向量，并將最終結(jié)果返回給output，這樣一來，輸入的單詞便成了詞向量。但這個操作只是詞嵌入的一部分，完整的詞嵌入還應(yīng)在詞嵌入中添加其它額外的信息，即：embedding_post_processor。
embedding_post_processor是詞嵌入操作必須添加進(jìn)去的第二部分信息，embedding_post_processor的輸入有：input_tensor、use_token_type、token_type_id、token_type_vocab_size,返回的特征向量將包含這些額外的信息，其維度和輸入單詞的詞向量一致。
6) 加入位置編碼特征
利用use_position_embedding 添加位置編碼信息。BERT 的Self_attention 中需要加入位置編碼信息，首先，利用full_position_embedding 初始化位置編碼，把每個單詞的位置編碼向量與詞嵌入向量相加，接著，根據(jù)當(dāng)前的序列長度做一個計算，如果序列長度為128，則對這128個位置進(jìn)行編碼。由于位置編碼僅包含位置信息，和句子的上下文語義無關(guān)，對于不同的輸入序列來說，雖然輸入序列的內(nèi)容各不相同，但是它們的位置編碼卻是相同的，所以位置編碼的結(jié)果向量和實(shí)際句子中傳的什么詞無關(guān)，無論傳的數(shù)據(jù)內(nèi)容是什么，它們的位置編碼均是一樣的。獲得位置編碼的輸出結(jié)果之后，在原詞嵌入輸出向量的基礎(chǔ)上，加上額外編碼獲得的特征向量和位置編碼向量，將三個向量求和，返回求和結(jié)果，到此為止，完成了BERT模型的輸入詞嵌入，得到了一個包含位置信息的詞向量，接下來，對這個向量進(jìn)行深入的操作。
圖 6
7) mask機(jī)制
在完成詞嵌入之后，接下來便是Transformer結(jié)構(gòu)了，在Transformer之前，先要對詞向量做一些轉(zhuǎn)換，即attention_mask ，創(chuàng)建一個mask矩陣：create_attention_mask_from_input_mask 。在前文提到的input_mask 中，只有mask=1 的詞參與到attention的計算當(dāng)中，現(xiàn)在需要把這個二維的mask轉(zhuǎn)換成為一個三維的mask，表示詞向量進(jìn)入attention的時候，哪幾個向量會參與到實(shí)際計算過程當(dāng)中。即在計算attention時，對輸入序列中128個單詞的哪些個單詞做attention計算，在這里，又額外地加入了一個mask處理操作。
圖 7
?
完成mask處理之后，接下來是構(gòu)建Transformer的Encode端，首先給Transformer傳入一些參數(shù)，如：input_tensor、attention_mask、hiden_size、head_num等等。這些參數(shù)在預(yù)訓(xùn)練過程中已經(jīng)設(shè)置好了，在進(jìn)行微調(diào)操作時，均不得對這些參數(shù)隨意更改。
在多頭attention機(jī)制中，每個頭生成一個特征向量，最終把各個頭生成的向量拼接在一起得到輸出的特征向量。
8) 構(gòu)建QKV 矩陣
接下來，是attention機(jī)制的實(shí)現(xiàn)，BERT 的attention機(jī)制是一個多層的架構(gòu)，在程序具體實(shí)現(xiàn)中，采用的是遍歷的操作，通過遍歷每一層，實(shí)現(xiàn)多層的堆疊?？偣残枰闅v12層，當(dāng)前層的輸入是前一層的輸出。attention機(jī)制中，有輸入兩個向量：from-tensor和to_tensor，而BERT 的attention機(jī)制采用的是self_attention，此時：
from-tensor=to_tensor=layer_input；
圖 8
在構(gòu)建attention_layer過程中，需要構(gòu)建K、Q、V 三個矩陣，K、Q、V矩陣是transformer中最為核心的部分。在構(gòu)建K、Q、V矩陣時，會用到以下幾個縮略字符：
-
B ?代表Batch Size ?即批大小 ?在這里的典型值設(shè)為 8；
-
F???代表 ?from-tensor??維度是128;
-
T???代表 to_tensor??維度是128;
-
N ? Number of Attention Head attention機(jī)制的頭數(shù)（多頭attention機(jī)制）在這里的典型值設(shè)為 12個頭；
-
H ? Size_per_head 代表每個頭中有多少個特征向量，在這里的典型值設(shè)為 64；
構(gòu)建Query 矩陣：構(gòu)建query_layer查詢矩陣，查詢矩陣由from-tensor構(gòu)建而來，在多頭attention機(jī)制中，有多少個attention頭，便生成多少個Query 矩陣，每個頭生成的Query 矩陣輸出對應(yīng)向量：
query_layer=【 B*F,N*H】， 即1024*768；
圖 9
構(gòu)建Key 矩陣: Key 矩陣由to-tensor構(gòu)建而來, 在多頭attention機(jī)制中，有多少個attention頭，便生成多少個Key 矩陣，每個頭生成的Key 矩陣輸出對應(yīng)向量：
key_layer=【 B*T,N*H】， 即1024*768；
圖 10
構(gòu)建Value矩陣: Value矩陣的構(gòu)建和Key 矩陣的構(gòu)建基本一樣,只不過描述的層面不同而已：
value_layer=【 B*T,N*H】， 即1024*768；
構(gòu)建QKV 矩陣完成之后,計算K矩陣和Q 矩陣的內(nèi)積，之后進(jìn)行一個Softmax操作。通過Value矩陣，幫助我們了解實(shí)際得到的特征是什么，Value矩陣和Key矩陣完全對應(yīng),維數(shù)一模一樣。
圖 11
9) 完成Transformer 模塊構(gòu)建
構(gòu)建QKV 矩陣完成之后,接下來，需要計算K矩陣和Q 矩陣的內(nèi)積，為了加速內(nèi)積的計算，在這里做了一個transpose轉(zhuǎn)換，目的是為了加速內(nèi)積的計算，并不影響后續(xù)的操作。計算好K矩陣和Q 矩陣的內(nèi)積之后，獲得了attention的分值：attention_score，最后需要利用Softmax操作將得到的attention的分值轉(zhuǎn)換成為一個概率：attention_prob。
在做Softmax操作之前，為了減少計算量，還需要加入attention_mask，將長度為128 的序列中不是實(shí)際有的單詞屏蔽掉，不讓它們參與到計算中來。在tensorflow中直接有現(xiàn)成的Softmax函數(shù)可以調(diào)用，把當(dāng)前所有的attention分值往Softmax里一傳，得到的結(jié)果便是一個概率值，這個概率值作為權(quán)重值，和Value矩陣結(jié)合在一起使用，即將attention_prob和Value矩陣進(jìn)行乘法運(yùn)算，便得到了上下文語義矩陣，即：
Context_layer=tf.matmul(attention_prob, value_layer);
圖 12
得到當(dāng)前層上下文語義矩陣輸出之后，這個輸出作為下一層的輸入，參與到下一層attention的計算中去，多層attention通過一個for循環(huán)的多次迭代來實(shí)現(xiàn)，有多少層attention（在這里是12層）就進(jìn)行多少層迭代計算。
10) 訓(xùn)練BERT 模型
做完self_attention之后，接下來是一個全連接層，在這里，需要把全連接層考慮進(jìn)來，利用tf.layer.dese 實(shí)現(xiàn)一個全連接層，最后要做一個殘差連接，注意：在全連接層的實(shí)現(xiàn)過程中，需要返回最終的結(jié)果，即將最后一層attention的輸出結(jié)果返回給BERT ，這便是整個Transformer 的結(jié)構(gòu)。
圖 13
總結(jié)一下上述整個過程，即Transformer 的實(shí)現(xiàn)主要分為兩大部分：第一部分是embedding 層，embedding 層將wordpiece詞嵌入加上額外特定信息和位置編碼信息，三者之和構(gòu)成embedding 層的輸出向量；第二部分是將embedding 層的輸出向量送入transformer結(jié)構(gòu)，通過構(gòu)建K、Q、V三種矩陣過，利用Softmax函數(shù)，得到上下文語義矩陣C , 上下文語義矩陣C不僅包含了輸入序列中各單詞的編碼特征,還包括了各單詞的位置編碼信息。
這就是BERT 模型的實(shí)現(xiàn)方式，理解了上述兩大部分的詳細(xì)過程，對BERT模型的理解便沒有什么太大問題了。以上十大步驟基本涵蓋了BERT 模型中的重要操作。
經(jīng)過BERT 模型之后，最終獲得的是一個特征向量，這個特征向量代表了最終結(jié)果。以上便是谷歌官方公布的開源MRPC 項(xiàng)目的全部過程。讀者在構(gòu)建自己特定任務(wù)的項(xiàng)目時候，需要修改的是如何將數(shù)據(jù)讀入BERT 模型的部分代碼，實(shí)現(xiàn)數(shù)據(jù)預(yù)處理。
