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

[TOC]

<a href="https://colab.research.google.com/github/davidho27941/ML_tutorial_notebook/blob/main/PyTorch_Lightning.ipynb" target="_parent\">
<img src="https://colab.research.google.com/assets/colab-badge.svg\" alt="Open In Colab\">
</a>

## 簡介

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

## 行前複習

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

> 也可瀏覽先前的[章節](https://hackmd.io/@davidho9713/HyfHBzcxY/%2FbGSXeP5iRpa8m0XwyircIw)來複習。

# 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`能帶給我們的好處

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

> 節錄自`PyTorch-Lightning`的[GitHub頁面](https://github.com/PyTorchLightning/pytorch-lightning)

## 開始使用`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`來完成對不同裝置的整合以及支援。

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

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

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

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

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

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

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

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

我們在`training_step`中利用`self.log`來對準確率以及損失值進行了紀錄，但我們在這裡不會進行介紹，有興趣的人可以先前往瀏覽相關[文件](https://pytorch-lightning.readthedocs.io/en/stable/extensions/logging.html#)來了解詳情。

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`來編寫的核心部份。
<center class="half">
<img src="https://i.imgur.com/0CuLKoh.png" width="300"><img src="https://i.imgur.com/NIE0QXT.png" width="300">
</center>

我們省去了`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有興趣的話，可以參考他們的[網站](https://pytorch-lightning.readthedocs.io/en/stable/_modules/pytorch_lightning/core/hooks.html)，這裡就不一一介紹了。

## 結語

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

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