基於PyTorch的「Keras」：除了核心邏輯通通都封裝

機器之心發表於2019-08-05
PyTorchKeras封裝

PyTorch Lightning 地址：https://github.com/williamFalcon/pytorch-lightning

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

看起來像 Keras 的 PyTorch

Keras 本身的目的就是對深度學習框架（TensorFlow、Theano）進行了進一步的 API 封裝。作為 TensorFlow 的高度封裝，Keras 的抽象層次非常高，很多 API 細節都隱藏了起來。雖然 PyTorch 同樣使用動態計算圖，也方便快捷，但總體上 Keras 隱藏的細節更多一些。

反觀 PyTorch，它提供一個相對較低階別的實驗環境，使使用者可以更加自由地編寫自定義層、檢視數值最佳化任務等等。例如在 PyTorch 1.0 中，編譯工具 torch.jit 就包含一種名為 Torch Script 的語言，它是 Python 的子語言，開發者使用它能進一步對模型進行最佳化。

用 PyTorch 寫模型，除了資料載入和模型定義部分外，整個訓練和驗證的邏輯、配置都需要我們手動完成，這些步驟都較為繁瑣。甚至可以說，研究者需要耗費相當多的精力處理這一部分的程式碼，還要祈禱不出 Bug。但是對於大多數研究實驗來說，訓練和驗證的迴圈體差不多都是一樣的，實現的功能也相當一致，所以為什麼不將這些通用的東西都打包在一起，這樣訓練不就簡單了麼？

William Falcon 正是這樣想的，他將 PyTorch 開發中的各種通用配置全都包裝起來，我們只需要寫核心邏輯就行。透過 PyTorch Lightning，PyTorch 就類似於 Keras，它能以更高階的形式快速搭建模型。

專案作者是誰

要完成這樣的工作，工作量肯定是非常大的，因為從超參搜尋、模型 Debug、分散式訓練、訓練和驗證的迴圈邏輯到模型日誌的列印，都需要寫一套通用的方案，確保各種任務都能用得上。所以 Facebook 的這位小哥哥 William Falcon 還是很厲害的。

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

他是一位 NYU 和 Facebook 的開發者。目前在 NYU 攻讀 PhD。從 GitHub 的活動來看，小哥是一位比較活躍的開發者。

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

這是一位披著 Keras 外衣的 PyTorch

Lightning 是 PyTorch 非常輕量級的包裝，研究者只需要寫最核心的訓練和驗證邏輯，其它過程都會自動完成。因此這就有點類似 Keras 那種高階包裝，它隱藏了絕大多數細節，只保留了最通俗易懂的介面。Lightning 能確保自動完成部分的正確性，對於核心訓練邏輯的提煉非常有優勢。

那麼我們為什麼要用 Lightning？

當我們開始構建新專案，最後你希望做的可能就是記錄訓練迴圈、多叢集訓練、float16 精度、提前終止、模型載入/儲存等等。這一系列過程可能需要花很多精力來解決各式各樣、千奇百怪的 Bug，因此很難把精力都放在研究的核心邏輯上。

透過使用 Lightning，這些部分都能保證是 Work 的，因此能抽出精力關注我們要研究的東西：資料、訓練、驗證邏輯。此外，我們完全不需要擔心使用多 GPU 加速會很難，因為 Lightning 會把這些東西都做好。

所以 Lightning 都能幫我們幹什麼？

下圖展示了構建一個機器學習模型都會經歷哪些過程，很多時候最困難的還不是寫模型，是各種配置與預處理過程。如下藍色的部分需要用 LightningModule 定義，而灰色部分 Lightning 可以自動完成。我們需要做的，差不多也就載入資料、定義模型、確定訓練和驗證過程。

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

下面的虛擬碼展示了大致需要定義的幾大模組，它們再加上模型架構定義就能成為完整的模型。

# what to do in the training loop
def training_step(self, data_batch, batch_nb):
# what to do in the validation loop
def validation_step(self, data_batch, batch_nb):
# how to aggregate validation_step outputs
def validation_end(self, outputs):
# and your dataloaders
def tng_dataloader():
def val_dataloader():
def test_dataloader():

除了需要定義的模組外，以下步驟均可透過 Lightning 自動完成。當然，每個模組可以單獨進行配置。

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

Lightning 怎麼用

Lightning 的使用也非常簡單，只需要兩步就能完成：定義 LightningModel；擬合訓練器。

以經典的 MNIST 影像識別為例，如下展示了 LightningModel 的示例。我們可以照常匯入 PyTorch 模組，但這次不是繼承 nn.Module，而是繼承 LightningModel。然後我們只需要照常寫 PyTorch 就行了，該呼叫函式還是繼續呼叫。這裡看上去似乎沒什麼不同，但注意方法名都是確定的，這樣才能利用 Lightning 的後續過程。

import os
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision.datasets 
import MNIST
import torchvision.transforms as transforms
import pytorch_lightning as ptl

class CoolModel(ptl.LightningModule):  
  def __init__(self):    
    super(CoolModel, self).__init__()    # not the best model...    
    self.l1 = torch.nn.Linear(28 * 28, 10)
  
  def forward(self, x):    
    return torch.relu(self.l1(x.view(x.size(0), -1)))
    
  def my_loss(self, y_hat, y):    
    return F.cross_entropy(y_hat, y)
      
  def training_step(self, batch, batch_nb):    
    x, y = batch    
    y_hat = self.forward(x)    
    return {'loss': self.my_loss(y_hat, y)} 

  def validation_step(self, batch, batch_nb):    
    x, y = batch    
    y_hat = self.forward(x)    
    return {'val_loss': self.my_loss(y_hat, y)}   
   
  def validation_end(self, outputs):   
    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()      
    return {'avg_val_loss': avg_loss}    
  def configure_optimizers(self):      
    return [torch.optim.Adam(self.parameters(), lr=0.02)]

@ptl.data_loader 
def tng_dataloader(self):   
  return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

@ptl.data_loader 
def val_dataloader(self):   
  return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

@ptl.data_loader 
def test_dataloader(self):   
  return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

隨後，第二步即擬合訓練器。這就比較類似 Keras 這類高階包裝了，它將訓練配置細節、迴圈體、以及日誌輸出等更加具體的資訊全都隱藏了，一個 fit() 方法就能自動搞定一切。

這相比以前寫 PyTorch 更加便捷精煉一些，而且分散式訓練也非常容易，只要給出裝置 id 就行了。

from pytorch_lightning import Trainer
from test_tube import Experiment

model = CoolModel()
exp = Experiment(save_dir=os.getcwd())

# train on cpu using only 10% of the data (for demo purposes)
trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)
# train on 4 gpus
# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3])

# train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)
# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)

# train (1 epoch only here for demo)
trainer.fit(model)

# view tensorflow logs 
print(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')
print('and going to http://localhost:6006 on your browser')

其他特性

Pytorch-Lightning 還可以和 TensorBoard 無縫對接。

基於PyTorch的「Keras」：除了核心邏輯通通都封裝

只需要定義執行的路徑：

from test_tube import Experiment
from pytorch-lightning import Trainer
exp = Experiment(save_dir = '/some/path')
trainer = Trainer(experiment = exp)

將 TensorBoard 連線到路徑上即可：

tensorboard -logdir /some/path

相關文章