PyTorch Lightning 入坑心得
前言
好吧我必須說這篇文寫得有點意義不明，純粹是記錄我工作某一天突然發現這東西之後的驚喜感，總之先前情提要一下 : 在一個風和日麗但下著暴雨的午後，那是被 tensorflow 2.0 migration 的各種眼花撩亂 API 氣爛而開始轉向 PyTorch 的一個多月，老實說 PyTorch 其實算是簡潔，但多半跟許多底層 API 一樣，對於純粹要進行研究時的開發者來說是個麻煩，舉幾個例子 :
- 手動添加 tensorboard 的程式碼 : tensorboard 是一個很方便的工具，尤其是網絡訓練複雜但你偏偏又需要監測一些時間序列指標如 loss 隨 epoch 的變化的時候，可以大幅度減低試錯時間，但在 pyTorch 底下就需要 customize 每一個步驟需要紀錄的時刻 (關鍵我每次都要思考一遍要不要用 self 把那個物件包在 class 中，不包覺得未來要搜索麻煩，要包起來又覺得多出很多不必要的空間浪費)
- 仍舊多得數不清的參數定義 : 其實神經網絡的構造本身就已經有很多參數需要定義了，比方說 hidden layer, embedding size, min_delta, max_epoch, all_dims, dropout, batch_size 等族繁不及備載，更別提通常從原始資料到進入模型前還有諸多更可怕的定義像是 negative sampling ratio, clamp range, sampling size 等，衍生的問題在於對於參數傳輸的不便，在一個巢狀定義或多重繼承的場景中，要測試多組參數別往往都得再另外寫一組程式去維護版本管理，或是花時間研究簡化方法等
- 重複性的內容定義 : pyTorch 雖說簡潔，但仍舊有不少地方需要重複造車，或是至少得多思考一步你會在哪些地方需要重複造車，而事先將其分隔。舉個例子，要將 tensor 移動到 GPU 上運算時，通常需要先定義好 device = ‘cuda:0’，並且針對每次 Dataloader 吐出來的向量用 .to(device=device) 的方式將其移動到 GPU 運算，而如果有特別輸出需求，往往又得移動回 CPU 才能轉化為 numpy 形式。此外，許多自定義的內容需要包含大量的 forward, backward 繼承方法，比如在 training 中定義一次，在 customized loss function 中也需要定義一次，一來一往還是會有不少麻煩。
- 許多參數傳輸層級的問題 : 這個最主要的點就在於 model 與 optimizer 兩個參數，在一個完整的 pipeline 中，我們必須時刻注意調用 zero_grad(), eval() 等方法來確保得出我們所需要的結果並且不會改動到 class 的 property
- 不支援 early stopping : 這個有夠煩，每次都得手動寫一遍甚麼 if loss ≥ thres : trigger += 1 啦有的沒的XD
當然，上面提及的內容其實有點雞蛋裡挑骨頭了，當你熟悉這個框架以後實際上寫程式碼的速度還是比起 tensorflow 要快一些，不過 tensorflow 還有開源 Keras 框架呀! 對於前期研究型的目標來說，仍然要方便一些，因此我當時一遍苦惱著 CUDA coding 的內容，一邊研究有沒有東西可以讓我更專注在處理更 high-level 的目標上，這時，噹噹噹噹，我發現了這東西 :
這是基於 pyTorch 而衍生出來的高級框架，老實說一般我在改框架之前心裡都還是有些猶豫，畢竟框架這東西雖說要學總是學得會，但畢竟時間成本擺在那，而且很多時候能否成功轉移還是仰賴你對底層框架的了解程度，不過老實說，我在學習過程中幾乎是沒費多少功夫，大概一個下午就能把原始程式碼改得七七八八，並且程式碼數量減少了 25% 左右，也不需要我為了避免混淆以及兼顧工廠原則定義一大堆亂七八糟的函數看了心累XD
因此，想來分享一篇關於 pyTorch Lightning 的幾個使用心得！希望也對正苦於研究型任務的你有所助益。
一、初心者上手 : 先來看看 Lightning 框架的基本步驟!
最簡單使用 pyTorch-Lightning 的方法就是使用這份文檔裡面所說的，Lightning in Two Steps !
首先，第 0 步我們當然得先把這個 package 給裝起來，在 command line 輸入 :
pip install pytorch-lightning或你是用 conda :
conda install pytorch-lightning -c conda-forge
安裝部分應該不用我多說了，接下來我們直接進入到第一步驟 :
Step 1: Define LightningModule
基本上，LightningModule 幾乎完全等價於 torch.nn.Module，因此你可以大膽地使用原先定義在裡頭的所有函式，這邊官方文檔以 AutoEncoder 為例，定義了以下的程式碼 :
import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as plclass LitAutoEncoder(pl.LightningModule):
def __init__(self):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(28*28, 64),
nn.ReLU(),
nn.Linear(64, 3)
)
self.decoder = nn.Sequential(
nn.Linear(3, 64),
nn.ReLU(),
nn.Linear(64, 28*28)
)
def forward(self, x):
# in lightning, forward defines the prediction/inference actions
embedding = self.encoder(x)
return embedding
def training_step(self, batch, batch_idx):
# training_step defined the train loop.
# It is independent of forward
x, y = batch
x = x.view(x.size(0), -1)
z = self.encoder(x)
x_hat = self.decoder(z)
loss = F.mse_loss(x_hat, x)
# Logging to TensorBoard by default
self.log('train_loss', loss)
return loss
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
return optimizer
這邊一個一個區塊講解，首先是在 __init__ 函式部分定義了兩串 nn.Sequential，也就是預先把模型網絡結構給定義好，這部分幾乎與原先的 nn.Module 沒有任何區別，而該模組透過 training step 以及 forward 來改寫掉以往的定義方式，forward 在 pytorch_lightning 裡頭是用來做 prediction 居多。過去光是在 training step 你一般需要這樣寫 :
import torch
from torch import nn
import torch.nn.functional as Ffrom tensorboardX import SummaryWriterclass LitAutoEncoder(nn.Module):
def __init__(self, **kwargs):
super().__init__()
self.encoder = nn.Sequential(
nn.Linear(28*28, 64),
nn.ReLU(),
nn.Linear(64, 3)
)
self.decoder = nn.Sequential(
nn.Linear(3, 64),
nn.ReLU(),
nn.Linear(64, 28*28)
) self.kwargs = kwargs def forward(self, x):
x = x.view(x.size(0), -1)
z = self.encoder(x)
x_hat = self.decoder(z)
return x def fit(self, optimizer, loader_train, max_epochs=20): writer = SummaryWriter('runs/exp-1')
for e in range(max_epochs):
for t, (x, y) in enumerate(loader_train):
x = x.to(device='cuda:0')
new_x = model(x)
loss = F.mse_loss(new_x, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
writer.add_scalar('loss_' + e, loss.item(), e)
看出來兩者的差別了嗎 ? 這裡來一一細數 :
- 省去了 Dataloader 吐資料的環節，前者統一將其用 batch 參數替代掉，並且會自動檢測現在可用的 CPU, GPU, TPU 資源，自動將 tensor 移動到裝置上進行運算。
- 將一些繁瑣的 optimizer.zero_grad(), loss.backward(), model.eval() 等都給包裝起來，你就不需要再重複寫一次這些東西，不要覺得這沒什麼，上面僅僅定義了 training step，這代表在傳統 nn.Module 框架中，你還得在 validation step 去考慮一下這些東西是否加上去 (比如說要加上 with torch.no_grad() 等)。當然，現在還是支持你自己定義這些參數，比如你可以 :
class LitAutoEncoder(pl.LightningModule):
def backward(self, loss, optimizer, optimizer_idx):
loss.backward()- 只要直接調用 self.log，會自動幫你將資訊儲存到 tensorboard 裡面，這大幅減少了需要使用的程式碼數量
- 把 epoch 的迭代步驟包裝到更高層級的接口(待會步驟二會說到)，這代表現在可以用更有效的方式去閱讀哪些步驟應該使用 epochs 參數，這件事情為甚麼這麼重要? 在於理想的情況下 epoch 或 iteration 應該只會是調參的一個環節，你不需要把它寫在底層函式裡頭，當然你可以預先在 __init__() 函數中就先定義這個參數，但往往這樣的方式又導致該參數與其他層級的參數 — 有些甚至極其無用的比如 verbose 混淆在一起，增加後續維護與閱讀障礙。
- 最後，終於不需要在 validation 的時候寫一次一模一樣的東西了，在 pytorch_lightning 中只要在該 Module 裡多定義一個函數，就能夠自動幫你執行 Validation 啦 (還附帶 tqdm 的進度條XD) :
def validation_step(self, batch, batch_idx):
loss = MSE_loss(...)
self.log('val_loss', loss)以上就是在第一步裡面給出的一些糖果，接下來我們來看第二步驟，也就是透過這個 Module 去執行一些如 training 等目標 :
Step 2: Fit with Lightning Trainer
這邊更簡單，你只需要這樣寫 :
# init model
autoencoder = LitAutoEncoder()
# most basic trainer, uses good defaults (auto-tensorboard, checkpoints, logs, and more)
# trainer = pl.Trainer(gpus=8) (if you have GPUs)
trainer = pl.Trainer()
trainer.fit(autoencoder, train_loader)在上方提及的一些如 epochs 參數等，其實都會在 Trainer 裡面出現，這樣的結構我個人認為增加了不少程式碼的可閱讀性，以及減輕了額外定義參數的需求 (當然前提是你真的不需要 customize 他們)
就這樣！兩步驟的快速上手到此告一段落，希望你可以體會到一些 Lightning 框架帶來的便利性，此外他們也做了個很可愛的 Gif 動畫來展示這些演變 :
二、往實務開發邁進 : 在 Lightning 裡面達成 OO 效果 !
一般在 pyTorch coding 中也不是如此簡單地把 Model 結構定義好就行，通常你還需要額外幾個步驟來讓整段 code 可以好好運行你的專案，而在 lightning 中當然也有照顧到這些需求，以下我們來一個一個說明 :
01. 自定義 Dataset
在 pyTorch 中，你可以透過繼承 torch.utils.data.Dataset 的方式來自定義你的 Dataset，並且透過 torch.utils.data.Dataloader 的物件來達到 batch loading, sampling 等效果。一般而言，假設我們在一個推薦系統的場景，我有一個資料集 Samples 長這樣 :
╔═════════╦═════════╦═══════╗
║ User_id ║ Item_id ║ Values║
╠═════════╬═════════╬═══════╣
║ AAA ║ DDD ║ 2 ║
║ BBB ║ EEE ║ 1 ║
║ CCC ║ FFF ║ 3 ║
╚═════════╩═════════╩═══════╝此外，我們同時有兩個預先處理好的 user embedding array 與 item embedding array，這時你可以用這樣的方式來撰寫你的自定義資料集。
在繼承 Dataset 的自定義 class 中，你只需要自定義 __getitem__ 的迭代方法，以及 __len__ 的資料集總長度，就可以透過 DataLoader 方式去批量地吐出資料，上面的程式碼是比較繁瑣地把標籤與值分為兩個資料集 (對我來說比較實用，畢竟全部放一起還要額外記住 index)，如果是同一個資料集的話，在 Dataloader 裡面也有辦法去做 sampler 去切割 train/test。
而這樣的方式會有甚麼麻煩呢 ? 主要在於資料的前處理以及 Dataloader 方法沒辦法合併成一個模組，這會造成程式碼東一塊西一塊，非常不方便，如果說前處理的需求龐大到你需要另外拉出來先解決並儲存還可以理解，但許多時候該前處理只是一些簡單的任務比方說負採樣，這時候分散開來就顯得有些雞肋。此外，split() 方法，是否導入 prediction 等你都還是要分開存放，這無形中也增加了許多不便。
另外一個不便的點是，有時候許多的前處理運算與資料加載是可以預先在 GPU/內存運行的，這可能會幫你節省後續整個模型訓練時間的10%~20%，但一般來說你還是需要花一點時間 customize code，將其移轉到 GPU 上面或預先放在內存中，具體可以參考 :
pytorch_lightning 在這裡給出的解決方案是另一個大雜燴的 class 定義 :
class MyDataModule(LightningDataModule):
def __init__(self):
super().__init__()
self.train_dims = None
self.vocab_size = 0
def prepare_data(self):
# called only on 1 GPU
download_dataset()
tokenize()
build_vocab()
def setup(self):
# called on every GPU
vocab = load_vocab()
self.vocab_size = len(vocab)
self.train, self.val, self.test = load_datasets()
self.train_dims = self.train.next_batch.size()
def train_dataloader(self):
transforms = ...
return DataLoader(self.train, batch_size=64)
def val_dataloader(self):
transforms = ...
return DataLoader(self.val, batch_size=64)
def test_dataloader(self):
transforms = ...
return DataLoader(self.test, batch_size=64)A datamodule encapsulates the five steps involved in data processing in PyTorch:
1. Download / tokenize / process.
2. Clean and (maybe) save to disk.
3. Load inside Dataset.
4. Apply transforms (rotate, tokenize, etc…).
5. Wrap inside a DataLoader.
此時可以看到，上面的 split、前處理等步驟全部被包在一個 class 之中處理完畢，且 prepare_data 與 setup 兩步驟已經預先將運算過程在GPU上呼叫，並盡可能將其保存到 DISK 中，能夠省下你一些時間。不過我個人對於目前的方案還是認為可以加強一下，其一當然是最初自定義 Dataset 的部分如果你的資料集比較複雜，終究你還是要自定義的，本質上沒有解決太多問題，此外雖然說將 train, val split 也包含在 Module 裡頭，但程式碼量並不會因此節省多少，主要的助益在於自動化保存到 disk，並且也提供一些簡單的 transform 方法可以對接。
但整體而言，在使用上還是十分令人愉快的，何況雖然每個人情況不一樣，我工作中的確得到了10~20%的運算加速。
02. 應用 Early Stopping
幾乎每個神經網絡的搭建都會包含 Early Stopping 功能去防止過擬合，但在傳統的 pyTorch 裡頭，幾乎需要自己重新手寫一遍，比如類似這樣 :
可以看到我幾乎得另外定義一個 early stopping 函式，並記錄每次 loss 的情況然後寫在迭代裡頭，只能說，非常不方便而且對於嘗試不同種類的 Early Stopping 而言並不是如此的友善 (比如說我就是想用 MAE 當停止條件，哪怕我是 train 在 BCE loss 上，怎樣XDD)。
在 pytorch_lightning 中，可以在 trainer 裡面指定 callbacks 參數來使用 early stopping，非常簡單 :
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
def validation_step(...):
self.log('val_loss', loss)early_stop_callback = EarlyStopping(
monitor='val_loss',
min_delta=0.00,
patience=3,
verbose=False,
mode='max'
)
trainer = Trainer(callbacks=[early_stop_callback])
通常在 validation step 中，本身就需要紀錄 loss 了，恰巧 lightning 透過 self.log 的方式幫助我們進行紀錄，因此在 callback 中，我們可以很簡單定義 monitor 參數來選擇我們要監測的 loss 對象。(這東西用過真的回不去XD)。
如果你要自定義，官方也有對應的說明 :
class MyEarlyStopping(EarlyStopping):
def on_validation_end(self, trainer, pl_module):
# override this to disable early stopping at the end of val loop
pass
def on_train_end(self, trainer, pl_module):
# instead, do it at the end of training loop
self._run_early_stopping_check(trainer, pl_module)03. Argument Parser
一般而言，在研發階段最方便迭代的方式就是有一些 Parser 的方法能夠讓你嘗試想要的參數，而最常用的方法當然是 Python 的標準函式庫 argparse，lightning 也可以支援該方法 :
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument('--layer_1_dim', type=int, default=128)
args = parser.parse_args()04. 部署與預測
Lightning 最方便的一個性質是他幾乎是完全等同於 nn.Module，這代表你在 nn.Module 裡面怎麼樣去做數據預測，在 Lightning 中就是一樣的方法。
一般來說，我們在 nn.Module 裡面都會考慮在訓練完一個模型後將其保存，nn.Module 裡面提供了許多種類的保存方法，舉例來說，官方推薦的作法是我們僅保存模型訓練後的參數集，而不需保存整個模型 — 畢竟後續的預測僅需要參數即可執行，很多時候不需要浪費空間儲存一整個模型。
#save model parameters
torch.save(model.state_dict(), PATH)#load model paramters
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
該方法會將 nn.Module 中的 state_dict 字典物件給保存下來，這裡面包含了神經網絡中的各層權重、偏誤值等等，儲存的物件格式則為 pickle 檔案。在加載模型預測的時候也可以直接調用 load_state_dict() 方法將參數給還原，當然，這種作法不僅僅在於預測，在設定 checkpoint 時也非常方便，你可以透過類似這樣的方式來設定中斷點 :
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, PATH)具體其他的酷炫功能請參考文檔 :
pytorch_lightning 則在該基礎上加入了自動化保存機制等新功能，一般來說，你在透過 Trainer 執行模型運算的時候，會自動在進度條裡面看到類似這樣的輸出 :
其他東西我們都很熟悉，包括輸出各層資訊(類型、參數數量)以及當前 Epoch 的訓練誤差等等，如果你仔細看上面的進度條，你會發現在最右邊有一個 v_num 參數，每一次執行的時候，該數字會自動增加，這代表的是當前執行的版本號碼，而 lightning 會自動根據該版本號碼去保存模型資訊，你可以透過在 Trainer 物件中帶參數 default_root_dir 的方式來保存。
trainer = Trainer(default_root_dir='/your/path/to/save/checkpoints')至於我們說過的 checkpoint，就更加方便了。與 EarlyStopping 一樣，lightning 透過在 Trainer 物件裡面指定一個 callbacks 參數即可 :
from pytorch_lightning.callbacks import ModelCheckpoint
# saves a file like: my/path/sample-mnist-epoch=02-val_loss=0.32.ckpt
checkpoint_callback = ModelCheckpoint(
monitor='val_loss',
dirpath='my/path/',
filename='sample-mnist-{epoch:02d}-{val_loss:.2f}',
save_top_k=3,
mode='min',
)
trainer = Trainer(callbacks=[checkpoint_callback])05. Transfer Learning
一樣，由於 lightning 與 nn.Module 幾無二致，如果我想用一個預先訓練好的模型去 warm start 其他的訓練任務，只需要 :
class Encoder(torch.nn.Module):
...
class AutoEncoder(pl.LightningModule):
def __init__(self):
self.encoder = Encoder()
self.decoder = Decoder()
class CIFAR10Classifier(pl.LightingModule):
def __init__(self):
# init the pretrained LightningModule
self.feature_extractor = AutoEncoder.load_from_checkpoint(PATH)
self.feature_extractor.freeze()
# the autoencoder outputs a 100-dim representation and CIFAR-10 has 10 classes
self.classifier = nn.Linear(100, 10)
def forward(self, x):
representations = self.feature_extractor(x)
x = self.classifier(representations)
...在上述的例子中，CIFAR10Classifier 在 __init__ 階段預先加載了 AutoEncoder 訓練好的模型參數，後續在 forward 裡面就可以直接應用來做預測任務，非常方便！
06. 自定義損失函數 Customized Loss Function
很多時候我們也需要自定義損失函數，然而可惜的是目前在 pytorch_lightning 中，我並沒有發現更多新奇的做法去處理這類問題，好在由於其與 nn.Module 的高度兼容性，你完全可以用一樣的方法寫一個自定義損失函數並直接在 lightning 物件裡面呼叫他，比如你可以自己寫一個帶權重的 BCE Loss Function :
按照上面的 training step，你只需要直接呼叫他就可以了 :
criterion = BCELoss(max_score=5, weight=[1, 2, 3, 4, 5])def training_step(self, batch, batch_idx):
# training_step defined the train loop.
# It is independent of forward
x, y = batch
x = x.view(x.size(0), -1)
z = self.encoder(x)
x_hat = self.decoder(z)
loss = criterion(x_hat, x)
# Logging to TensorBoard by default
self.log('train_loss', loss)
return loss
三、前人種樹後人乘涼
當然，其他還有一些新奇功能幫助你更快開展你的專案，在一個叫做 Lightning-Bolts 的庫中，存放著許多 SOTA 的預訓練模型、自定義資料集與損失函數等等，這邊來簡單介紹一些 :
01. 我不需要深度學習，怎麼辦?
沒問題，即使你只需要很簡單的部署一個邏輯回歸，也不需要你大費周章做一個一層 Full-Connected 帶 Sigmoid 的神經網絡，Bolt 裡面有支援與 sklearn 的完美對接，通過 SklearnDataModule，你就可以將 Sklearn 原先需要做成 X, y 的形式也包在 loader 裡面 :
from pl_bolts.models.regression import LogisticRegression
from pl_bolts.datamodules import SklearnDataModule
from sklearn.datasets import load_boston
import pytorch_lightning as pl
# sklearn dataset
X, y = load_boston(return_X_y=True)
loaders = SklearnDataModule(X, y)
model = LogisticRegression(input_dim=13)
# try with gpus=4!
# trainer = pl.Trainer(gpus=4)
trainer = pl.Trainer()
trainer.fit(model, loaders.train_dataloader(), loaders.val_dataloader())
trainer.test(test_dataloaders=loaders.test_dataloader())02. 想要直接從前人的 SOTA 研究中開始
沒問題，舉個例來說，我現在想用 Google 在去年公布的 SimCLR 來做自監督學習，Bolt 裡面也完美幫你處理了許多類似的 SOTA 模型 :
from pl_bolts.models.self_supervised import SimCLR
from pl_bolts.models.self_supervised.simclr.transforms import SimCLRTrainDataTransform, SimCLREvalDataTransform
import pytorch_lightning as pl
# data
train_data = DataLoader(MyDataset(transforms=SimCLRTrainDataTransform(input_height=32)))
val_data = DataLoader(MyDataset(transforms=SimCLREvalDataTransform(input_height=32)))
# model
weight_path = 'https://pl-bolts-weights.s3.us-east-2.amazonaws.com/simclr/simclr-cifar10-v1-exp12_87_52/epoch%3D960.ckpt'
simclr = SimCLR.load_from_checkpoint(weight_path, strict=False)
simclr.freeze()由上述的例子中你可以看到，現在你可以直接從 pl_bolts.models.self_supervised 中來找到預訓練好的 SimCLR 並加載他的 checkpoint，你的模型就可以直接從人家處理好的結果中開始迭代，是不是感覺特別舒服?
當然，實務中能夠直接抓取別人的模型使用的情況少之又少，這包含了資料格式的不同、訓練目標的不同等多種問題，但不可否認的是，直接先從一個已經跑好的結果開始迭代總是要快得多一些，能夠免除你一些很奇怪的試錯成本，比如若簡單資料在預訓練模型上都不理想的時候，恭喜你可以直接回歸到 Data Quality 開始檢查了，省去你重頭搭建框架的時間。
03. 長得很奇怪的損失函數與資料集
有一些比較麻煩的情況在於，如果你不幸需要把模型搞得很複雜，比方說 self-supervised learning 的情況，可能傳統的損失函數或資料集建構方法也並不完全適用，比如上面提到的 SimCLR 就採用了一個神奇的 Normalized Temperature-Scaled Cross Entropy Loss，長這樣 :
Well，好吧，像這種時候你就無可避免得要自己手寫 Loss Function 了，不過在 bolt 裡面，這些問題也解決掉啦！因為人家都幫你寫好了XD 直接從 pl_bolts.losses 或 pl_bolts.datasets 去 import 即可。
結語
好吧寫這篇純粹是我個人的私心，主要是某天發現 lightning 之後用起來覺得愛不釋手，上手起來也並不複雜，因此就寫了一篇文章來安麗我的發現XD
平心而論，新的架構對於程式碼的節省也沒有到這麼誇張，沒辦法動輒省下你50%的工作量等等，但簡潔的呈現與壁壘分明的參數傳遞還是讓人覺得十分舒適的，真心推推哈！
我也在這個 Publication 裡面開了一個亂寫的 Section，希望之後還能分享更多我發現的酷玩意~
