HackMD
利用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來編寫的核心部份。
- The image file may be corrupted
- The server hosting the image is unavailable
- The image path is incorrect
- The image format is not supported
- The image file may be corrupted
- The server hosting the image is unavailable
- The image path is incorrect
- The image format is not supported
我們省去了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的功能來一步一步讓整個流程更加的有彈性以及組織。
