1. 程式人生 > 實用技巧 >TransCoder程式碼詳解(三):DAE/BT的訓練過程

TransCoder程式碼詳解(三):DAE/BT的訓練過程

前言

ATP的上一篇blog裡講了這個模型是怎麼用Masked Language Model進行預訓練的。其實就跟BERT的方法一樣。

在預訓練完了以後,模型的路徑下肯定就有了一個預訓練好的checkpoint。

接下來的步驟就應該是把這個預訓練好的模型load進來,在它的基礎上再來進行DAE和BT的訓練。

DAE和BT的訓練主要目的是訓練模型的decoder。

建立模型build_model

因為這次是要建立一個完整的transformer,所以這裡呼叫TransformerModel類的建構函式兩次,分別例項化一個encoder和一個decoder。通過is_encoder這個引數來控制當前建立的到底是哪一部分。

注意到build_model函式的第二句涉及到一個separate_decoders引數。這個引數在train.py的getparser裡面有定義,意思是“是否給不同的語言用不同的decoder”。換句話說,如果這個引數設定為true,現在訓練涉及到三種語言,那麼就會有三個decoder。然而readme裡面給的命令列中,這個引數的值是false。所以原模型只有一個decoder。

接下來載入已經用MLM預訓練好的引數即可。

關於decoder引數的初始化,ATP沒有特別搞懂這部分程式碼。由於原文中說這部分基本上是照搬XLM的訓練過程,ATP看了一下其他人對於XLM模型的解析。

這篇部落格認為,在XLM模型中,是先用CLM/MLM預訓練encoder,然後用這個encoder的引數來初始化後續MT訓練過程要用到的encoder和decoder。實際上decoder的大部分也是使用之前預訓練的引數來初始化的,只不過decoder比encoder多了一個enc-dec attention部分,這部分引數沒法與encoder對應,所以也就沒有辦法初始化,需要在後續訓練中進行學習。

意思就是如下圖所示,紅色框框的部分是隨機初始化的,而藍色和綠色框框的部分初始是用同一批引數填進去的。

ATP認為這說法非常可信,因為原論文中也是這個意思。

Denoising auto-encoding 訓練過程 mt_step

從字面意思上也可以看出,這個DAE的訓練實際上就是一個翻譯的過程,只不過源語言和目標語言是同一種。

在一番初始化過後,首先需要注意到的是開頭generate batch的部分。這一部分的作用是取出一批資料用於訓練。

# generate batch
if lang1 == lang2:
    (x1, len1) = self.get_batch('ae', lang1)
    (x2, len2) = (x1, len1)
    (x1, len1) = self.add_noise(x1, len1)
else:
    (x1, len1, _, _), (x2, len2, _, _) = self.get_batch('mt', lang1, lang2)

在DAE的過程中,lang1和lang2是相等的。於是它執行if後面的三個語句。可以發現,它的源資料和目標資料實際上是同一批資料,只不過源資料加了一些噪聲。

加噪聲的過程add_noise也是定義在trainer的類裡面的,包括三個步驟:打亂順序、丟棄字元與打mask。

資料處理好以後就是非常經典的步驟,先送encoder,再送decoder,計算loss用來優化,就結束了。

Back-translation 訓練過程 bt_step

bt_step是整個訓練過程裡面比較特殊的一個步驟。

這個bt_step過程有三個關於語言的引數lang1,lang2和lang3。其中lang1和lang3是相同的,lang2和lang1是不同的,恰好對應了back-translation的過程:先將A翻譯成B,再將B翻譯回A。於是整個bt_step也順理成章地分成了兩個大的部分。

bt_step函式的大致流程如下所示。

def bt_step(self, lang1, lang2, lang3, lambda_coeff, sample_temperature):
    """
    Back-translation step for machine translation.
    """
    # 初始化各種引數以及A->B過程要用的“臨時”模型
    ......

    if self.generation_index >= len(self.generated_data):
        # 取資料
        ......
            
        # 開始生成A->B的資料
        with torch.no_grad():

            # 切換到eval模式
            self.eval_mode()

            # 將初始資料放入臨時的encoder和decoder,生成資料
            ......

    # 整理資料準備B->A的過程使用
    ......

    # 切換到train模式,開始訓練B->A過程
    self.train_mode()
    
    # 將資料送入encoder和decoder
    ......

    # 計算loss並優化
    ......

    # 收尾工作
    ......

這個過程中,實際上可以看作整個過程只有一個encoder和一個decoder。但在生成A->B的資料時,從當前encoder和decoder臨時拷貝了一份引數相同的模型來用。

由於整個A->B的過程使用的是pytorch模型的eval模式,並且在torch.no_grad()的作用範圍下,這份拷貝出來的模型並不會參與訓練,只是負責生成一批資料。

A->B的過程輸出的結果,會作為B->A過程的輸入。B->A的過程與上文提到的普通MT過程差別不大。

B->A過程用到的ground_truth就是A->B過程的輸入,也就是未經處理的原始資料。這兩者共同計算出loss以後就可以優化。

由於只有在B->A的過程中將模型切換到了pytorch的train模式,所以優化的目標只限於B->A過程涉及到的部分。