Pytorch

Pytorch Lightning 入門筆記

Posted by MingLun Allen Wu on Thursday, April 16, 2020

前言

在了解 Pytorch 的基本使用後，透過 Pytorch Lightning 框架能夠讓我們更有效率的進行開發。如果對於 Pytorch 的基本使用還不熟悉的讀者，可以先看看我先前寫的文章:

從零開始 - Pytorch入門懶人包

簡介

透過 Pytorch 撰寫 Deep Learning 相關程式碼時，程式碼大致可分成兩種類型:

  1. 依照專案有所不同: 這類程式碼會根據開發的需求而改變，例如「資料的前處理」、「模型的架構」。

  2. 跨專案重複使用: 這類程式碼在每一個應用中都會存在，且相似性非常高，例如「將資料送入 DataLoader」、「計算 Loss」、「透過 Optimizer 更新模型的 weight」。

Pytorch Lightning 希望替使用者簡化的就是「跨專案重複使用」這部分的程式碼，讓使用者能省下精力及時間去處理更為重要的核心模型部分。

你可以將其視為是一種「樣式指南」(類似 Python 的 PEP8)，透過一定規則將訓練過程中的幾個重要步驟封裝起來(Train, 計算Loss, Optimizer, 更新參數)。

你可能很好奇這樣做的用意是什麼？ 你將可以透過特別的 Trainer Class 來進行操作，不再需要花費時間撰寫 Evaluation 的 Loss、或是在每一個 Epoch 結束時顯示 Loss、固定幾個 Step 要顯示當前訓練狀態、儲存 Validation表現最好的 Checkpoint。 只要你先按照 Pytorch Lightning 的格式封裝程式碼，在訓練時的瑣碎細節它會幫你處理好。

下面這張圖則是 Pytorch Lightning 官網的例子，用來說明 Pytorch Lightning 想要簡化的部分。

我們來看個較為實際的例子:

model = SomeModel() # 假設這是你按照Pytorch Lightning封裝好的Model.
# 透過GPU訓練
trainer = Trainer(gpus=1, precision=16)
trainer.fit(model) 

# 透過CPU訓練
trainer = Trainer(gpus=0)
trainer.fit(model)

# 進行Evaluation
trainer.test()

從上面的範例中可以看到，原先冗長的訓練過程現在只需要透過短短的幾行程式碼即可實現。將訓練時的細節按照特定規則封裝好後，就可直接透過 Trainer Class 來進行訓練。

一、安裝

透過pip即可進行安裝:

pip install pytorch-lightning==0.7.1

二、基本概念

Pytorch Lightning 將深度學習的程式碼分成三種類型:

  1. Research Code: 整個應用的核心架構，可能會根據任務的內容進行調整、或是在開發過程中加入自己的想法。通常可能會包含幾個核心元件:

    • 模型架構定義
    • Train/Val/Test資料切分
    • Optimizer定義
    • Train/Val/Test Step Computation.

  2. Engineer Code:

    「EarlyStopping」, 「將資料送入GPU」等常見且不同專案使用方式都類似的程式碼。這部分的程式碼是 Pytorch Lightning 最希望能簡化的。

  3. Non-Essential Code:

    檢查梯度、視覺化等偏向輔助性質的功能。

Pytorch Lightning 希望使用者只需要定義 Research Code，而 Engineer Code 由它來代為處理、Non-Essential 則是根據使用者的需要自行選用，由於不會影響正常使用，接下來我們著重探討如何重新將Research Code 組織為 Pytorch Lightning 的格式。

三、 Research Code

針對模型開發中最為重要的就是 Research Code了，在 Pytorch Lightning 中透過 pl.LightningModule 模組來實現，這個模組繼承了 torch.nn.Module 的功能，所以使用上其實跟使用原生 Pytorch 是相當類似的。

實作時只需在定義的 Class 上繼承 pl.LightningModule 即可。

import pytorch_lightning as pl

class YourOwnNet(pl.LightningModule):

    def somemethod(self):
        pass

在繼承 pl.LightingModule 後，有幾項重要的 Method 必須被 Override，未來 Pytorch Lightning 才能自動地進行訓練，以下將這些方法及目的條列出來:

  • prepare_data():

    負責資料的載入(包含 Training Set, Evaluation Set, Test)都撰寫在方法中。 自動訓練時會優先執行此 Method 以獲取資料。

  • configure_optimizer(self):

    • 回傳特定的優化器 (torch.nn.optimizer)
    • 如果對於優化的參數有任何設定(例如只想對特定參數進行調整)，也是在這個地方調整。
    • 在自動訓練時，會呼叫此方法來取得 Optimizer。
      def configure_optimizers(self):
    return Adam(self.parameters(), lr=1e-3)

  • train_dataloader(self) / val_dataloader(self) / test_dataloader(self):

    • 回傳 Pytorch 的 DataLoader Object，

    • 這三個方法定義了 training/ validation/ test 時使用的 DataLoader。其中會使用到的資料可以透過 prepare_data() 先行取得。

    • 在自動訓練的過程中，會呼叫這邊定義的 Function 來取得不同時期(Train/Validation/Test)的 DataLoader。

      def train_dataloader(self):
    # self.train_dataset 通常由 self.prepare_data()產生.
    return DataLoader(self.train_dataset, batch_size= self.batch_size, shuffle=True)

  • forward(self, args):

    • 定義當前這個模型 forward propagation 時所要進行的動作。

    • 特別注意在建立 instance 時，直接將參數傳入 instance 等同於呼叫 forward function。來個具體例子。

      model = TestModel()
x = torch.Tensor([1, 2, 3])
output1 = model(x)
output2 = model.forward(x)
output1 == output2  # 實際上這兩個是相同的.

  • training_step/val_step/test_step(self, batch, batch_idx):

    • 分別定義了 traing/ validation/ test 時的每一個 Step 所要進行的任務，其中較特別的是參數batch, batch_idx，這是會自動從 train/val/test_dataloader() 中的 DataLoader 取出一個 batch。具體使用方式請見下方範例。

    • 此類型方法的回傳值有一定的格式，需要回傳一個 Python 的 Dictionary，其中包含:

      • loss: 當前 Step 計算出來的 Training Loss. 用於接下來的 Backward Propagation 及參數調校。

      • log: 包含當前訓練的狀況，在訓練時會自動回傳成進度條。 如下圖所示: 

        def training_step(self, batch, batch_idx):
    x, y = batch  # 這裡的batch會對應到 self.train_dataloader()所回傳的DataLoader Object，對其取用一個batch.

    output = self.forward(x) # 將Input進行 forward propagation.
    criterion = nn.CrossEntropyLoss()
    loss = criterion(output, y)
    logs = {"loss": loss}
    return {"loss": loss, "log": logs}

  • (training/validation/test)_epoch_end(self, outputs):

    • 定義了「不同時期訓練完一個epoch時需要執行的事項」，舉例來說可能會在 Validation 的一個 Epoch 結束時計算所有 Step 平均的 Loss 是多少。

    • outputs 參數為「當前 Epoch 中所有 Step 的{“loss”: loss, “logs”: logs}」(也就是每一個 self.training_step() 的回傳值)。

    • 以下範例示範了在「每一個 Validation Epoch 結束後計算 Average Loss」。

      def validation_epoch_end(self, outputs):
    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
    avg_val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()

    tensorboard_logs = {'val_loss': avg_loss, 'avg_val_acc': avg_val_acc}
    return {'avg_val_loss': avg_loss, 'progress_bar': tensorboard_logs}

  • 統整版本:

class PytorchLightningModel(pl.LightningModule): # 這邊一定要繼承pl.LightningModule
    def __init__(self): # 初始化時可以將基本設定傳入。
        super().__init__()
        self.ln1 = nn.Linear(768, 256)
        self.ln2 = nn.Linear(256, 3)

    def prepare_data(self): # 此方法會在初始化後優先執行。 所以可以在此方法中先將會用到的資料都讀取進來.
        self.train_set = read_data("train") # read_data是自定義的讀取資料Method. 可以按照自己需求調整
        self.test_set = read_data("test")
        self.val_set = read_data("validation")
        logging.info("===== Data is ready... =====")
    
    def configure_optimizer(self): # 自動訓練時會呼叫此方法來獲取Optimizer.
        return Adam(self.parameters(), lr=1e-3) # 這邊注意要調整的參數是`self.parameters()`
    
    # 以下三個方法則是設定進行訓練及驗證時所要使用的Data Loader格式。
    def train_dataloader(self):
        return DataLoader(self.train_set, batch_size= self.batch_size, shuffle=True) 
    
    def val_dataloader(self):
        return DataLoader(self.val_set, batch_size= self.batch_size, shuffle=True) 
    
    def test_dataloader(self):
        return DataLoader(self.test_set, batch_size= self.batch_size, shuffle=True) 
    
    def forward(self, x): # 定義模型在forward propagation時如何進行.
        output = self.ln1(x)
        output = self.ln2(x)
        return output
    
    def training_step(self, batch, batch_idx): # 定義訓練過程的Step要如何進行
        x, y = batch # 從self.train_dataloader()的Data Loader取一個batch出來。
        output = self.forward(x)
        criterion = nn.CrossEntropyLoss()
        loss = criterion(output, y)
        logs = {'loss': loss}
        return {'loss':loss, 'log':logs}
    
    def validation_step(self, batch, batch_idx): # 定義Validation如何進行，以這邊為例就再加上了計算Acc.
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y)

        # acc
        a, y_hat = torch.max(logits, dim=1)
        val_acc = accuracy_score(y_hat.cpu(), y.cpu())
        val_acc = torch.tensor(val_acc)

        return {'val_loss': loss, 'val_acc': val_acc}
    
    def validation_epoch_end(self, outputs): # 在Validation的一個Epoch結束後，計算平均的Loss及Acc.
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()

        tensorboard_logs = {'val_loss': avg_loss, 'avg_val_acc': avg_val_acc}
        return {'avg_val_loss': avg_loss, 'progress_bar': tensorboard_logs}
    
    def test_step(self, batch, batch_idx): #定義 Test 階段
        x, y = batch
        logits = self.forward(x)
        loss = F.cross_entropy(logits, y)

        # acc
        a, y_hat = torch.max(logits, dim=1)
        val_acc = accuracy_score(y_hat.cpu(), y.cpu())
        val_acc = torch.tensor(val_acc)

        return {'test_acc': val_acc}

四、透過Trainer自動訓練

Training

在你重新將程式碼改寫為 Pytorch Lightning 格式後，輕鬆的部分來了！接下來我們只需要透過 Trainer Class即可自動處理訓練。我們看看以下的範例:

model = PytorchLightningModel(param1=768, param2=5) # 自行封裝成Pytorch Lightning的模型

# Trainer 有不同的參數可以調整訓練時的行為
trainer = pl.Trainer() # 使用CPU
trainer = pl.Trainer(gpus=1) # 使用GPU
trainer = pl.Trainer(fast_dev_run=True) #訓練時，使用單一個batch作為ㄧ個Epoch，目的是快速的確認當前的模型設置有無結構上的問題(快速地跑完Train -> Validation -> Test)。

trainer.fit(model) # 呼叫.fit() 就會自動進行Model的training step及validation step.

原先從 CPU 切換為 GPU，可能需要針對所有的變數、模型進行裝置的變更。但現在只需要調整 Trainer 的 gpus 參數即可。

而訓練時如果需要以 dry-run 的方式，確定 Model 從頭到尾的結構沒有問題，現在也只需要設定 fast_dev_run 即可。

另外 Trainer Class 也針對了「使用多張GPU」、「accumulate_grad_batches」提供參數進行調整，細節的部分可以參考官方網站。 連結在此

Evaluation / Inference

在訓練完模型後，你也能夠載入先前訓練好的 Checkpoint File 並且重新進行 Evaluation/Inference.

model = PytorchLightningModel(param1=768, param2=5) # 自行封裝成Pytorch Lightning的模型
trainer = Trainer(resume_from_checkpoint="PATH", gpus=1)
trainer.test(model)

五、後記

從原生 Pytorch 轉移到 Pytorch Lightning 並不需要花費過多成本來學習新的語法，其概念只是將程式碼「重新組織」，將對應的程式碼放置到特定的 Method 中，如此一來在進行訓練時，將能省下許多工程面的實作時間、程式碼也會變得簡潔許多。

See Also

進職場學到的Git - (1) Gitlab CI/CD
NLP深度學習不能不用的套件 - Transformers
不再被大量超參數及模型表現淹沒！使用MLFlow進行實驗管理
CATALOG
    FEATURED TAGS
    ansible book concept fast-api fast-api-tutorial music note pytest-101 python test thought tool wsgi