天天看點

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

目前,無論是從性能、結構還是業界應用上,Transformer 都有很多無可比拟的優勢。本文将介紹 PaddlePaddle 的 Transformer 項目,我們從項目使用到源碼解析帶你玩一玩 NMT。隻需千行模型代碼,Transformer 實作帶回家。

其實 PyTorch、TensorFlow 等主流架構都有 Transformer 的實作,但如果我們需要将它們應用到産品中,還是需要修改很多。

例如谷歌大腦建構的 Tensor2Tensor,它最開始是為了實作 Transformer,後來擴充到了各種任務。對于基于 Tensor2Tensor 實作翻譯任務的使用者,他們需要在 10 萬+行 TensorFlow 代碼找到需要的部分。

PaddlePaddle 提供的 Transformer 實作,項目代碼隻有 2000+行,簡潔優雅。如果我們使用大 Batch Size,那麼在預測速度上,PaddlePaddle 複現的模型比 TensorFlow 官方使用 Tensor2Tensor 實作的模型還要快 4 倍。

項目位址:

https://github.com/PaddlePaddle/models/tree/develop/fluid/PaddleNLP/neural_machine_translation/transformer

1. Transformer 怎麼用

相比此前 Seq2Seq 模型中廣泛使用的循環神經網絡,Transformer 使用深層注意力機制獲得了更好的效果,目前大多數神經機器翻譯模型都采用了這一網絡結構。此外,不論是新興的預訓練語言模型,還是問答或句法分析,Transformer 都展現出強大的模組化能力。

相比傳統 NMT 使用循環層或卷積層抽取文本資訊,Transformer 使用自注意力網絡抽取并表征這些資訊,下圖對比了不同層級的特點: 

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

不同網絡的主要性質,其中 n 表示序列長度、d 為隐向量次元、k 為卷積核大小。例如單層計算複雜度,一般句子長度 n 都小于隐向量次元 d,那麼自注意力層級的計算複雜度最小。

如上所示,Transformer 使用的自注意力模型主要擁有以下優點,1)網絡結構的計算複雜度最低;2)由于序列操作數複雜度低,模型的并行度很高;3)最大路徑長度小,能夠更好地表示長距離依賴關系;4)模型更容易訓練。

現在,如果我們需要訓練一個 Transformer,那麼最好的方法是什麼?當然是直接跑已複現的模型了,下面我們将跑一跑 PaddlePaddle 實作的 Transformer。

1.1 處理資料

在 PaddlePaddle 的複現中,百度采用原論文測試的 WMT'16 EN-DE 資料集,它是一個中等規模的資料集。這裡比較友善的是,百度将資料下載下傳和預處理等過程都放到了 gen_data.sh 腳本中,包括 Tokenize 和 BPE 編碼。

在這個項目中,我們既可以通過腳本預處理資料,也可以使用百度預處理好的資料集。首先最簡單的方式是直接運作 gen_data.sh 腳本,運作後可以生成 gen_data 檔案夾,該檔案夾主要包含以下檔案:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

其中 wmt16_ende_data_bpe 檔案夾包含最終使用的英德翻譯資料。

如果我們從頭下載下傳并預處理資料,那麼大概需要花 1 到 2 個小時完成預處理。為此,百度也提供了預處理好的 WMT'16 EN-DE 資料集,它包含訓練、驗證和測試所需要的 BPE 資料和字典。

其中,BPE 政策會把稀疏詞拆分為高頻的子詞,這樣既能解決低頻詞無法訓練的問題,也能合理降低詞表規模。

如果不采用 BPE 的政策,要麼詞表的規模變得很大,進而使訓練速度變慢或者顯存太小而無法訓練;要麼一些低頻詞會當作未登入詞處理,進而得不到訓練。

預處理資料位址:

https://transformer-res.bj.bcebos.com/wmt16_ende_data_bpe_clean.tar.gz

如果我們有其它資料集,例如中英翻譯資料,也可以根據特定的格式進行定義。例如用空格分隔不同的 token(對于中文而言需要提前用分詞工具進行分詞),用\t 分隔源語言與目智語句對。

1.2 訓練模型

如果需要執行模型訓練,我們也可以直接運作訓練主函數 train.py。如下簡要配置了資料路徑以及各種模型參數:

# 顯存使用的比例,顯存不足可适當增大,最大為1
export FLAGS_fraction_of_gpu_memory_to_use=0.8
# 顯存清理的門檻值,顯存不足可适當減小,最小為0,為負數時不啟用
export FLAGS_eager_delete_tensor_gb=0.7
python -u train.py \
  --src_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --trg_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --special_token '<s>' '<e>' '<unk>' \
  --train_file_pattern gen_data/wmt16_ende_data_bpe/train.tok.clean.bpe.32000.en-de \
  --token_delimiter ' ' \
  --use_token_batch True \
  --batch_size 1600 \
  --sort_type pool \
  --pool_size 200000 \
  n_head 8 \
  n_layer 4 \
  d_model 512 \
  d_inner_hid 1024 \
  prepostprocess_dropout 0.3      

此外,如果顯存不夠大,那麼我們可以将 Batch Size 減小一點。為了快速測試訓練效果,我們将模型調得比 Base Transformer 還小(降低網絡的層數、head 的數量、以及隐層的大小)。

上面僅展示了小部分的超參設定,更多的配置可以在 GitHub 項目 config.py 檔案中找到。預設情況下,模型每疊代一萬次儲存一次模型,每個 epoch 結束後也會儲存一次 cheekpoint。此外,在我們訓練的過程中,預設每一百次疊代會列印一次模型資訊,其中 ppl 表示的是困惑度,困惑度越小模型效果越好。

在單機訓練中,預設使用所有 GPU,可以通過 CUDA_VISIBLE_DEVICES 環境變量來設定使用的 GPU,例如 CUDA_VISIBLE_DEVICES='0,1',表示使用 0 号和 1 号卡進行訓練。

1.3 預測推斷

訓練完 Transformer 後就可以執行推斷了,我們需要運作對應的推斷檔案 infer.py。我們也可以在推斷過程中配置超參數,但注意超參需要和前面訓練時保持一緻。

python -u infer.py \
  --src_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --trg_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \
  --special_token '<s>' '<e>' '<unk>' \
  --test_file_pattern gen_data/wmt16_ende_data_bpe/newstest2016.tok.bpe.32000.en-de \
  --token_delimiter ' ' \
  --batch_size 32 \
  model_path trained_models/iter_100000.infer.model \
  n_head 8 \
  n_layer 4 \
  d_model 512 \
  d_inner_hid 1024 \
  prepostprocess_dropout 0.3
  beam_size 5 \
  max_out_len 255      

相比模型的訓練,推斷過程需要一些額外的超參數,例如配置 model_path 指定模型所在目錄、設定 beam_size 和 max_out_len 來指定 Beam Search 每一步候選詞的個數和最大翻譯長度。這些超參數也可以在 config.py 中找到,該檔案對這些超參都有注釋說明。

執行以上預測指令會将翻譯結果直接打出來,每行輸出是對應行輸入得分最高的翻譯。對于使用 BPE 的英德資料,預測出的翻譯結果也将是 BPE 表示的資料,是以需要還原成原始資料才能進行正确評估。如下指令可以将 predict.txt 内的翻譯結果(BPE 表示)恢複到 predict.tok.txt 檔案中(tokenize 後的資料):

sed -r 's/(@@ )|(@@ ?$)//g' predict.txt > predict.tok.txt      

在未使用內建方法的情況下,百度表示 base model 和 big model 在收斂後,測試集的 BLEU 值參考如下:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

這兩個預訓練模型也提供了下載下傳位址:

2. Transformer 怎麼改

如果我們想要訓練自己的 Transformer,那麼又該怎樣了解并修改 PaddlePaddle 代碼呢?如果我們需要根據自己的資料集和任務改代碼,除了前面資料預處理過程,模型結構等子產品有時也需要修改。這就需要我們先了解源代碼了,PaddlePaddle 的源代碼基本都是基礎的函數或運算,我們很容易了解并使用。

對于 PaddlePaddle 不熟悉的讀者可查閱文檔,也可以看看入門教程,了解基本編寫模式後就可以看懂整個實作了。

PaddlePaddle 官網位址:

http://paddlepaddle.org/paddle

如 Seq2Seq 一樣,原版 Transformer 也采用了編碼器-解碼器架構,但它們會使用多個 Multi-Head 注意力、前饋網絡、層級歸一化和殘差連接配接等。下圖從左到右展示了原論文所提出的 Transformer 架構、Multi-Head 注意力和标量點乘注意力。

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

上圖右邊的點乘注意力就是标準 Seq2Seq 模型中的注意力機制,中間的 Multi-head 注意力其實就是将一個隐層資訊切分為多份,并單獨計算注意力資訊,使得一個詞與其它多個目标詞的注意力資訊計算更精确。最左邊為 Transformer 的整體架構,編碼器與解碼器由多個類似的子產品組成,後面将簡要介紹這些子產品與對應的 PaddlePaddle 代碼。

2.1 點乘注意力

注意力機制目前在機器翻譯中已經極其流行了,我們可以認為 Transformer 是一種堆疊多層注意力網絡的模型,它采用的是一種名為經縮放的點乘注意力機制。這種注意力機制使用經縮放的點乘作為作為評分函數,進而評估各隐藏狀态對目前預測的重要性,如下是該注意力的表達式:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

其中 Query 向量與 (Key, Value ) 向量在 NMT 中相當于目智語輸入序列與源語輸入序列,Query 與 Key 向量的點乘,經過 SoftMax 函數後可得出一組歸一化的機率。這些機率相當于給源語輸入序列做權重平均,即表示在生成新的隐層資訊的時候需要關注哪些詞。

在 Transformer 的 PaddlePaddle 實作中,經縮放的點乘注意力是在 Multi-head 注意力函數下實作的,如下所示為上述表達式的實作代碼:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

在這個函數中,q、k、v 和公式中的一樣,attn_bias 用于 Mask 掉標明的特定位置(encode 的 self attention 和 decoder 端的 encode attention 都是屏蔽掉 padding 的詞;decoder 的 self attention 屏蔽掉目前詞後面的詞,目的是為了和解碼的過程保持一緻),是以在給不同輸入權重時忽略該位置的輸入。

如上 product 計算的是 q 和 k 之間的點乘,且經過根号下 d_key(key 的次元)的縮放。這裡我們可以發現參數 alpha 可以直接對矩陣乘法的結果進行縮放,預設情況下它為 1.0,即不進行縮放。在 Transformer 原論文中,作者表示如果 d_key 比較小,那麼直接點乘和帶縮放的點乘差别不大,是以他們認為高維情況下可能不帶縮放的乘積太大而令 Softmax 函數飽和。

weights 表示對輸入的不同元素權重,即不同輸入對目前預測的重要性,訓練中也可以對該權重進行 Dropout。最後 out 表示按照 weights 對輸入 V 進行權重和,得出來就是目前注意力的運算結果。

2.2 Muti-head 注意力

Multi-head 注意力其實就是多個點乘注意力并行地處理并最後将結果拼接在一起。一般而言,我們可以對三個輸入矩陣 Q、V、K 分别進行線性變換,然後分别将它們投入 h 個點乘注意力函數并拼接所有的輸出結果。

這種注意力允許模型聯合關注不同位置的不同表征子空間資訊,我們可以了解為在參數不共享的情況下,多次執行點乘注意力。如下所示為 Muti-head 注意力的表達式:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

其中每一個 head 都為一個點乘注意力,不同 head 的輸入是相同 Q、K、V 的不同線性變換。

總體而言,PaddlePaddle 的 Multi-head 注意力實作分為幾個步驟:先為 Q、K、V 執行線性變換;再變換次元以計算點乘注意力;最後計算各 head 的注意力輸出并合并在一起。

2.2.1 線性變換

如前公式所示,Muti-head 首先要執行線性變換,進而令不同的 head 關注不同表征空間的資訊。這種線性變換即乘上不同的權重矩陣,且模型在訓練過程中可以學習和更新這些權重矩陣。在如下的 PaddlePaddle 代碼中,我們可以直接調用全連接配接層 layers.fc() 完成線性變換。

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

直接調用全連接配接層會自動為輸入建立權重,且我們要求不使用偏置項和激活函數。這裡比較友善的是,PaddlePaddle 的 layers.fc() 函數可以接受高維輸入,省略了手動展平輸入向量的操作。是以這裡有 num_flatten_dims=2,即将前兩個次元展平為一個次元,第三個次元保持不變。

例如對于輸入張量 q 而言,線性變換的輸出次元應該是 [batch_size,max_sequence_length,d_key * n_head],最後一個次元即 n_head 個 d_key 維的 Query 向量。每一個 d_key 維的向量都會饋送到不同的 head,并最後拼接起來。

2.2.2 次元變換

為了進行 Multi-Head 的運算,我們需要将線性變換的結果進行 reshape 和轉置操作。現在我們将這幾個張量的最後一個次元分割成不同的 head,并做轉置以便于後續運算。

具體而言,輸入張量 q、k 和 v 的次元資訊為 [bs, max_sequence_length, n_head * hidden_dim],我們希望把它們轉換為 [bs, n_head, max_sequence_length, hidden_dim]。

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

如上使用 layers.reshape() 和 layers.transpose() 函數完成分割與轉置。其中 layers.reshape() 在接收輸入張量後會按照形狀 [0, 0, n_head, d_key] 進行轉換,其中 0 表示從輸入張量對應維數複制出來。此外,因為 inplace 設定為 True,那麼 reshape 操作就不會進行資料的複制,進而提升運算效率。

後面的轉置就比較簡單了,隻需要按照次元索引将第「1」個次元和第「2」個次元交換就行了。此外為了更快地執行推斷,PaddlePaddle 實作代碼還做了非常多的優化,例如這部分後續會對推斷過程的緩存和處理流程進行優化。

2.2.3 合并

前面已經介紹過點乘注意力了,那麼上面對 q、k、v 執行次元變換後就可直接傳入點乘注意力函數,并計算出 head_1、head_2 等注意力結果。現在最後一步隻需要将這些 head 拼接起來就完成了整個過程,也就完成了上面 Multi-head 注意力的計算式。

因為每一個批量、head 和時間步都會計算得出一個注意力向量,是以總體上注意力計算結果的次元資訊為 [bs, n_head, max_sequence_length, hidden_dim]。如果要将不同的 head 拼接在一起,即将 head 這個次元合并到 hidden_dim 中去,是以合并的過程和前面次元變換的過程正好相反。

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

如上合并過程會先檢驗次元資訊,然後先轉置再 reshape 合并不同的 head。注意在原論文中,合并不同的 head 後,還需要再做一個線性變換,這個線性變換的結果就是 Muti-head 注意力的輸出了。

最後,我們再将上面的四部分串起來就是 Transformer 最核心的 Multi-head 注意力。了解了各個子產品後,下面串起來就能愉快地看懂整個過程了:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

當然,如果編碼器和解碼器輸入到 Multi-head 注意力的 q 與 (k、v) 是相同的,那麼它又可稱為自注意力網絡。

2.3 前饋網絡

對于每一個編碼器和解碼器子產品,除了殘差連接配接與層級歸一化外,重要的就是堆疊 Muti-head 注意力和前饋網絡(FFN)。前面我們已經解決了 Multi-head 注意力,現在需要了解主位置的前饋網絡了。直覺而言,FFN 的作用是整合 Multi-head 注意力生成的上下文向量,是以能更好地利用從源語句子和目智語句子抽取的深度資訊。

如下所示在原論文中,前饋網絡的計算過程可以表達為以下方程:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

前饋網絡的結構很簡單,一個 ReLU 激活函數加兩次線性變換就完成了。如下基本上隻需要調用 PaddlePaddle 的 layers.fc() 就可以了:

想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下
想要千行代碼搞定Transformer?這份高效的PaddlePaddle官方實作請收下

現在基本上核心操作就定義完了,後面還有更多子產品與架構,例如怎樣利用核心操作搭建編碼器子產品與解碼器子產品、如何搭建整體 Transformer 模型等,讀者可繼續閱讀原項目中的簡潔代碼。整體而言,包括上面代碼在内,千行代碼就可以完全弄懂 Transformer,PaddlePaddle 的 Transformer 複現值得我們仔細讀一讀。

此外,在這千行模型代碼中,為了給訓練和推斷加速,還有很多特殊技巧。例如在 Decoder 中加入對 Encoder 計算結果的緩存等。加上這些技巧,PaddlePaddle 的實作才能在大 Batch Size 下實作 4 倍推斷加速。

因為本身 PaddlePaddle 代碼就已經非常精煉,通過它們也很容易了解這些技巧。基本上看函數名稱就能知道大緻的作用,再結合文檔使用就能完全讀懂了。

最後,除了模型架構,整個項目還會有其它組成部分,例如訓練、推斷、資料預處理等等。這些代碼同樣非常簡潔,我們可以根據實際需求閱讀并修改它們。總體而言,PaddlePaddle 的 Transformer 實作确實非常适合了解與修改。想要跑一跑神經機器翻譯的同學,PaddlePaddle 的 Transformer 實作确實值得推薦。

本文為機器之心原創,轉載請聯系本公衆号獲得授權。

繼續閱讀