利用PyTorch-Lightning最佳化訓練流程

利用PyTorch-Lightning最佳化訓練流程

The image file may be corrupted
The server hosting the image is unavailable
The image path is incorrect
The image format is not supported

簡介

在先前的章節中，我們瞭解了如何利用PyTorch建立資料流、模型以及進行訓練的流程。我們可以體會在在PyTorch中建立訓練所需要的所有元素是相對複雜的（與TensorFlow相比）。然而，我們可以透過PyTorch-Lightning這個模組所提供的各種功能來補足PyTorch所缺少的部份，並得到超越TensorFlow的彈性以及可塑性。在這個章節，我們將會初步介紹PyTorch-Lightning所能帶來的便利性，以及與過往使用原生PyTorch之間的差異。

行前複習

在開始了解如何利用PyTorch-Lightning函式庫來最佳化我們的訓練流程之前，我們先來複習如何在利用原生PyTorch函式庫來建構訓練的流程。

也可瀏覽先前的章節來複習。










































































































# Import libraries
import torchvision 
import torch 
import numpy as np
import pandas as pd

# Transform
# Step 1: Convert to Tensor
# Step 2: Normalize with mean = 0.5 and std= 0.5
transform = torchvision.transforms.Compose(
    [torchvision.transforms.ToTensor(),
     torchvision.transforms.Normalize((0.5,), (0.5,)),]
)

# Obtain MNIST training dataset.
# Transform is a custom functon define by user, which allow user to apply the transformation when loading dataset.
mnist_train = torchvision.datasets.MNIST(root='MNIST', 
                            download=True, 
                            train=True, 
                            transform=transform)
                            
# Obtain MNIST test dataset
mnist_test = torchvision.datasets.MNIST(root='MNIST', 
                            download=True, 
                            train=False, 
                            transform=transform)

# Build train data loader
trainLoader = torch.utils.data.DataLoader(mnist_train, 
                                          batch_size=64, 
                                          shuffle=True, 
                                          pin_memory=True, 
                                          num_workers=4)

# Build test data loader
testLoader = torch.utils.data.DataLoader(mnist_test, 
                                         batch_size=64, 
                                         shuffle=False, 
                                         pin_memory=True, 
                                         num_workers=4)

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

model = torch.nn.Sequential(
            torch.nn.Linear(in_features=784, out_features=128),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=128, out_features=64),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=64, out_features=10),
            torch.nn.LogSoftmax(dim=1)
            ).to(device)

EPOCHS = 10 
LR = 1e-2 
OPTIMIZER = 'adam' 

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)

for epoch in range(EPOCHS):
    # 定義變數用以儲存單次訓練週期中的loss以及acc值
    running_loss = list()
    running_acc = 0.0

    for times, data in enumerate(trainLoader):

        # 宣告模型訓練狀態
        model.train()

        # 取得資料，並將其傳遞至相應裝置
        inputs, labels = data[0].to(device), data[1].to(device)

        # 將影像維度改為第一維度長度為inputs.shape[0]，第二維度為所剩維度展平的長度
        inputs = inputs.view(inputs.shape[0], -1)

        # 將既存梯度歸零
        optimizer.zero_grad()

        # 將資料導入模型，並取得模型輸出
        outputs = model(inputs)

        # 將模型輸出轉為整數
        predicted = torch.max(outputs.data, 1)[1]
        
        # 利用損失函數計算loss
        loss = criterion(outputs, labels)

        # 進行反向傳播
        loss.backward()

        # 更新參數
        optimizer.step()

        # 紀錄週期內的損失值
        running_loss.append(loss.item())

        # 計算週期內正確預測的數量
        running_acc += (labels==predicted).sum().item()
    
    # 計算週期為單位的loss以及acc
    _epoch_loss = torch.tensor(running_loss).mean()
    _epoch_acc = running_acc/(len(trainLoader)*64)

    # 輸出資訊
    print(f"Epoch : {epoch+1}, Epoch loss: {_epoch_loss:.4f}, Epoch Acc: {_epoch_acc:.2f}")
print('Training Finished.')

想必對於大部分的人來說，這種需要對每一個步驟都進行設定的作法，相較下是比較令人感到麻煩的。在TensorFlow中都不需要親力親為的部份，在PyTorch中卻需要一項一項的處理，會花費掉開發者許多寶貴的時間。這時候就是PyTorch-Lightning函式庫大顯身手的好時機了。

關於利用`PyTorch-Lightning`能帶給我們的好處

不在需要處理變數與硬體之間的關係（可以把.to()丟掉了）。
因為大多數的工程部份被抽象化了，所以程式碼變得更加容易閱讀。
更容易重現結果。
因為少了棘手的工程部份，所以減少了失誤率。
因為保有PyTorch的原生功能，所以能在最大化彈性的情景下刪除大部分樣版程式。
可與大部分機器學習工具整合。
具有比PyTorch更好的運行效率。

節錄自PyTorch-Lightning的GitHub頁面

開始使用`PyTorch-Lightning`

在瞭解了PyTorch-Lightning能帶給我們怎樣的好處之後，我們可以開始來動手試試看將劑有訓練流程改寫為基於PyTorch-Lightning的形式。為此，首先我們需要安裝PyTorch-Lightning函式庫。

安裝PyTorch-Lightning函式庫

因為PyTorch-Lightning有發布在Pypi上，所以我們可以利用pip install pytorch-lightning來安裝PyTorch-Lightning函式庫。

pip3 install --quiet pytorch-lightning

利用`PyTorch-Lightning`建立訓練流程

利用PyTorch-Lightning建立訓練流程的過程很簡單，簡單來說就是「將必要的部份補齊」就好了。PyTorch-Lightning提供了各種已經設定好的類別，使用者只需要透過子類化(Sub-Classing)的方式繼承PyTorch-Lightning所建立好的類別，並將必要的部份補上/改寫即可。接下來就讓我們來看看如何實作。

匯入必要函式庫並設定部份參數















import os
import torch
from pytorch_lightning import LightningModule, Trainer
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchmetrics import Accuracy
from torchvision import transforms
from torchvision.datasets import MNIST
import pytorch_lightning as pl
import torchmetrics

!rm -rf lightning_logs/
AVAIL_GPUS = min(1, torch.cuda.device_count())
BATCH_SIZE = 256 if AVAIL_GPUS else 64

建立基於Lightning模組的機器學習模型物件

基於PyTorch-Lightning的功能，我們可以將過往的好幾個步驟整合成一個物件，並透過pl.Trainer來完成對不同裝置的整合以及支援。

在過往的流程中，我們必須要依照以下流程來建立整個模型以及資料管線：

建立DataLoader以及相關ETL管線。
建立模型。
設定超參數、優化器損失函數等等必要的部份。
利用for迴圈開始訓練並進行訓練、反向傳播、計算損失值等等操作。

在導入了PyTorch-Lightning之後，以上繁瑣的的部份都不再需要。你只需要在繼承了pl.LightningModule的物件中設定以下部份：

設定模型以及向前傳播的流程。(在物件中建立training_step以及validatation_step方法)
設定資料該如何被導入。 (於trainer.fit()中導入或是在物件中建立prepare_data或xxxx_dataloader方法)
設定優化器。 (在物件中建立configure_optimizers方法)

就可以了。在準備好物件之後，只需要透過pl.Trainer以及trainer.fit()函數的協助，就能透無痛的將你的訓練流程佈署在GPU、TPU以及其他可協助加速運算的硬體設備上進行訓練以及推論了。

以下是一個利用繼承了pl.LightningModule的物件來建立的MNIST模型物件。其包含了以下幾個元素：

__init__：整個物件的基本設定。
forward：模型如何向前傳播。
training_step：定義訓練的運算流程。
configure_optimizers：定義所使用的優化器。

在此之上，若有需要也可以自行添加validation_step來定義驗證的運算流程。

我們在training_step中利用self.log來對準確率以及損失值進行了紀錄，但我們在這裡不會進行介紹，有興趣的人可以先前往瀏覽相關文件來了解詳情。



















class MNISTModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)
        self.accuracy = torchmetrics.Accuracy()
    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        out = self(x)
        self.accuracy(out, y)
        loss = F.cross_entropy(self(x), y)
        self.log("train_loss", loss, on_epoch=True)
        self.log("train_acc", self.accuracy, on_epoch=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

在設定好模型物件後，我們就可以開始設定pl.Trainer()並利用它來進行訓練了。

進行訓練

在進行訓練之前，我們必須先建立一個pl.Trainer物件，並告知我們所希望使用的硬體加速設備、要跑幾個週期進行訓練，以及其他的設定。在建立pl.Trainer物件後，我們可以利用trainer.fit()搭配模型以及相關參數來開始我們的訓練。












dataset = MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor())
train, val = random_split(dataset, [55000, 5000])

mnist_model = MNISTModel()
trainer = pl.Trainer(max_epochs=10, 
                     gpus=AVAIL_GPUS, 
                     progress_bar_refresh_rate=20,
                     )

trainer.fit(mnist_model, 
            DataLoader(train, num_workers=4, pin_memory=True, batch_size=BATCH_SIZE), 
            DataLoader(val, num_workers=4, pin_memory=True, batch_size=BATCH_SIZE))

利用TensorBoard了解訓練的成果

因為我們上方有利用pl.LightningModule的self.log功能對準確率以及損失值進行紀錄，所以我們可以透過TensorBoard來了解經由這樣的訓練流程設計，我們是否能成功對的對模型進行訓練。


%load_ext tensorboard
%tensorboard --logdir lightning_logs/

我們可以發現，PyTorch-Lightning很順利的依照我們所設定的流程進行了訓練。

小結

藉由PyTorch-Lightning的幫助，我們省下了許多不必要且重複的的工程部份，以及節省了很大篇幅的程式編寫，並且很大的程度提昇了整個程式碼的可讀性。下方的組圖中，左方是舊有的訓練流程，而又方是藉由PyTorch-Lightning來編寫的核心部份。

我們省去了to()、backward以及step等不必要的部份，並獲得了專心在設計模型以及調整架構的時間以及機會，並且也無須在安排變數應該在哪一個裝置上費心思。這些就是藉由導入PyTorch-Lightning所帶來的好處。

`PyTorch-Lightning`機制介紹

在前面的部份，我們介紹了如何利用PyTorch-Lightning來建立訓練流程。接下來我們就來詳細了解PyTorch-Lightning是如何運作，又是如何以上方的機制來簡化流程的建立的。

設計核心

PyTorch-Lightning的設核心為：以基於PyTorch的架構為基礎，建構一個具有高度彈性以及可擴充性的函式庫。PyTorch-Lightning所提供的所有主要函式，都是基於PyTorch原本就有的功能。以LightningModule為例，LightningModule是基於nn.Module來建構的，也就是在原生PyTorch中我們透過模型子類化建立模型時，必須繼承的對象，所以過往所有的模型設計可以無縫接軌直接改寫成使用LightningModule來建立。

基於「Hook」，但不僅止於「Hook」

PyTorch-Lightning的架構之所有能有極高的彈性，就在於其設計上均以「Hook」的方式來將各個功能擴充到主要的功能之中，讓使用者不管在建立模型、建立數據管線、建立驗證以及測試邏輯、甚至是設計預測函式，都能夠被包裝在一個主要的物件之中，不需要在建立額外的物件。若要舉例的話，就類似於tf.keras.Callback所提供的on_epoch_end以及on_epoch_start等在自定義回測函數時，所用到的各個部份。使用者只需要將繼承對象原有的函式進行覆寫，就能夠隨心所欲的自訂一所想要的功能。

LightningModule ― 一即為全，全即唯一

LightningModule是在PyTorch-Lightning的架構中，最核心的一部分。無論是建立模型、設定超參數、建立數據管線、設定優化器以及處理訓練狀態等，只要是能想到的功能都能往繼承了LightningModule的物件裡面塞，最後在外面套一個pl.Trainer()結束這一回合。可以說是要多優雅就有多優雅，也不需要管變數要放在哪一個硬體上，更不用管資料要怎要在分佈式系統上處理，基本上交給LightningModule加上pl.Trainer()的組合就沒問題。

怎麼做到的？因為所有在PyTorch-Lightning中的Hook以及延伸功能都與LightningModule有掛勾，所以才能提供這樣極具彈性的功能。無論是訓練及測試用的xxxx_step也好，處理數據的xxxx_dataloader也罷。都與LightningModule緊緊相連。使用者只需要將必要的部份改寫成自己的設計，剩下的都可以交由PyTorch-Lightning所預先設計好的架構來處理。

對PyTorch-Lightning所提供的Hook有興趣的話，可以參考他們的網站，這裡就不一一介紹了。

結語

在本章節中，我們介紹了如何利用PyTorch-Lightning所帶來的強大功能來最佳化我們的各項流程，也介紹了PyTorch-Lightning的核心設計理念以及架構的精隨所在。接下來我們會介紹如何利用PyTorch-Lightning的功能來一步一步讓整個流程更加的有彈性以及組織。

利用PyTorch-Lightning最佳化訓練流程

簡介

行前複習

關於利用PyTorch-Lightning能帶給我們的好處

開始使用PyTorch-Lightning

安裝PyTorch-Lightning函式庫

利用PyTorch-Lightning建立訓練流程

匯入必要函式庫並設定部份參數

建立基於Lightning模組的機器學習模型物件

進行訓練

利用TensorBoard了解訓練的成果

小結

PyTorch-Lightning機制介紹

設計核心

基於「Hook」，但不僅止於「Hook」

LightningModule ― 一即為全，全即唯一

結語

tags: Machine Learning Notebook 技術隨筆 機器學習 Python PyTorch PyTorch-Lightning

Read more

資料流

技術隨筆

機器學習

基於容器的MLflow和模型版本控制

關於利用`PyTorch-Lightning`能帶給我們的好處

開始使用`PyTorch-Lightning`

利用`PyTorch-Lightning`建立訓練流程

`PyTorch-Lightning`機制介紹

tags: `Machine Learning` `Notebook` `技術隨筆` `機器學習` `Python` `PyTorch` `PyTorch-Lightning`