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

pythontab發表於2019-04-04

(導語)想要做個神經機器翻譯模型?想要做個強大的Transformer?搞定這千行PaddlePaddle程式碼你也可以。

  目前,無論是從效能、結構還是業界應用上,Transformer 都有很多無可比擬的優勢。本文將介紹Paddle Paddle 的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

  Transformer怎麼用

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

  相比傳統NMT使用迴圈層或卷積層抽取文字資訊,Transformer使用自注意力網路抽取並表徵這些資訊,下圖對比了不同層級的特點:

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

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

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

  處理資料

  在Paddle的復現中,百度採用原論文測試的WMT'16 EN-DE 資料集,它是一箇中等規模的資料集。這裡比較方便的是,百度將資料下載和預處理等過程都放到了gen_data.sh指令碼中,包括Tokenize 和 BPE 編碼。

  在這個專案中,我們既可以透過指令碼預處理資料,也可以使用百度預處理好的資料集。首先最簡單的方式是直接執行gen_data.sh指令碼,執行後可以生成gen_data資料夾,該資料夾主要包含以下檔案:

  其中 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分隔源語言與目標語句對。

  訓練模型

  如果需要執行模型訓練,我們也可以直接執行訓練主函式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 '' '' '' \

  --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號卡進行訓練。

  預測推斷

  訓練完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 '' '' '' \

  --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 值參考如下:

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

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

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

  Transformer 怎麼改

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

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

  PaddlePaddle官網地址:http://paddlepaddle.org/paddle

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

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

  點乘注意力

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

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

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

  def scaled_dot_product_attention(q, k, v, attn_bias, d_key, dropout_rate):

  """

  Scaled Dot-Product Attention

  """

  product = layers.matmul(x=q, y=k, transpose_y=True, alpha=d_key**-0.5)

  if attn_bias:

  product += attn_bias

  weights = layers.softmax(product)

  if dropout_rate:

  weights = layers.dropout(

  weights,

  dropout_prob=dropout_rate,

  seed=ModelHyperParams.dropout_seed,

  is_test=False)

  out = layers.matmul(weights, v)

  return out

  在這個函式中,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進行加權和,得出來就是當前注意力的運算結果。

  Muti-head 注意力

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

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

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

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

  1.線性變換

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

  def __compute_qkv(queries, keys, values, n_head, d_key, d_value):

  """

  Add linear projection to queries, keys, and values.

  """

  q = layers.fc(input=queries,

  size=d_key * n_head,

  bias_attr=False,

  num_flatten_dims=2)

  # For encoder-decoder attention in inference, insert the ops and vars

  # into global block to use as cache among beam search.

  fc_layer = wrap_layer_with_block(

  layers.fc, fluid.default_main_program().current_block()

  .parent_idx) if cache is not None and static_kv else layers.fc

  k = fc_layer(

  input=keys,

  size=d_key * n_head,

  bias_attr=False,

  num_flatten_dims=2)

  v = fc_layer(

  input=values,

  size=d_value * n_head,

  bias_attr=False,

  num_flatten_dims=2)

  return q, k, v

  直接呼叫全連線層會自動為輸入建立權重,且我們要求不使用偏置項和啟用函式。這裡比較方便的是,Paddle 的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.維度變換

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

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

  def __split_heads_qkv(queries, keys, values, n_head, d_key, d_value):

  reshaped_q = layers.reshape(

  x=queries, shape=[0, 0, n_head, d_key], inplace=True)

  q = layers.transpose(x=reshaped_q, perm=[0, 2, 1, 3])

  reshape_layer = wrap_layer_with_block(

  layers.reshape,

  fluid.default_main_program().current_block()

  .parent_idx) if cache is not None and static_kv else layers.reshape

  transpose_layer = wrap_layer_with_block(

  layers.transpose,

  fluid.default_main_program().current_block().

  parent_idx) if cache is not None and static_kv else layers.transpose

  reshaped_k = reshape_layer(

  x=keys, shape=[0, 0, n_head, d_key], inplace=True)

  k = transpose_layer(x=reshaped_k, perm=[0, 2, 1, 3])

  reshaped_v = reshape_layer(

  x=values, shape=[0, 0, n_head, d_value], inplace=True)

  v = transpose_layer(x=reshaped_v, perm=[0, 2, 1, 3])

  return q, k, v

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

  後面的轉置就比較簡單了,只需要按照維度索引將第“1”個維度和第“2”個維度交換就行了。此外為了更快地執行推斷,Paddle實現程式碼還做了非常多的最佳化,例如這部分後續會對推斷過程的快取和處理流程進行最佳化。

  3.合併

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

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

  def __combine_heads(x):

  """

  Transpose and then reshape the last two dimensions of inpunt tensor x

  so that it becomes one dimension, which is reverse to __split_heads.

  """

  if len(x.shape) != 4:

  raise ValueError("Input(x) should be a 4-D Tensor.")

  trans_x = layers.transpose(x, perm=[0, 2, 1, 3])

  # The value 0 in shape attr means copying the corresponding dimension

  # size of the input as the output dimension size.

  return layers.reshape(

  x=trans_x,

  shape=[0, 0, trans_x.shape[2] * trans_x.shape[3]],

  inplace=True)

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

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

  q, k, v = __compute_qkv(queries, keys, values, n_head, d_key, d_value)

  q, k, v = __split_heads_qkv(q, k, v, n_head, d_key, d_value)

  ctx_multiheads = scaled_dot_product_attention(q, k, v, attn_bias, d_model,

  dropout_rate)

  out = __combine_heads(ctx_multiheads)

  # Project back to the model size.

  proj_out = layers.fc(input=out,

  size=d_model,

  bias_attr=False,

  num_flatten_dims=2)

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

  前饋網路

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

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

  前饋網路的結構很簡單,一個ReLU啟用函式加兩次線性變換就完成了。如下基本上只需要呼叫Paddle的layers.fc() 就可以了:

  def positionwise_feed_forward(x, d_inner_hid, d_hid, dropout_rate):

  hidden = layers.fc(input=x,

  size=d_inner_hid,

  num_flatten_dims=2,

  act="relu")

  if dropout_rate:

  hidden = layers.dropout(

  hidden,

  dropout_prob=dropout_rate,

  seed=ModelHyperParams.dropout_seed,

  is_test=False)

  out = layers.fc(input=hidden, size=d_hid, num_flatten_dims=2)

  return out

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

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

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

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


相關文章