在PyTorch上用”Keras”，分布式訓(xùn)練開箱即用，告別沒(méi)完沒(méi)了的Debug

魚羊 2019-08-07 08:59:49 來(lái)源：量子位

魚羊 發(fā)自 凹非寺

量子位 報(bào)道 | 公眾號(hào) QbitAI

在開始一個(gè)新的機(jī)器學(xué)習(xí)項(xiàng)目時(shí)，難免要重新編寫訓(xùn)練循環(huán)，加載模型，分布式訓(xùn)練……然后在Debug的深淵里看著時(shí)間嘩嘩流逝，而自己離項(xiàng)目核心還有十萬(wàn)八千里。

雖然這世上已經(jīng)有了神器Keras，能用幾條語(yǔ)句就輕松組建一個(gè)神經(jīng)網(wǎng)絡(luò)，但一想到它是站在Tensorflow的肩膀上，就讓人不禁想起江湖中的那句傳說(shuō)：

PyTorch 真香！

那么為什么不做一個(gè)PyTorch上的Keras呢？

來(lái)自Facebook的Willian Falcon小哥決定一試，他搞了個(gè)包裝器，把PyTorch中的各種通用配置全部包裝在一起。

這個(gè)PyTorch輕量級(jí)包裝器，就是PyTorch Lightning。

有了這樣一個(gè)快速研究框架，使用者只需關(guān)注核心訓(xùn)練和驗(yàn)證邏輯，繁瑣的工程細(xì)節(jié)通通自動(dòng)化一鍵完成，既能保證核心訓(xùn)練邏輯的正確性，又能保證最佳的實(shí)踐體驗(yàn)。

像閃電一樣迅疾

所以，Lightning到底有多好用？

在這張圖中，灰色部分代表Lightning能自動(dòng)完成的部分，而藍(lán)色的部分則能夠根據(jù)使用者的需求，被定義成任意底層模型，可以是你自己的新模型，也可以是預(yù)訓(xùn)練模型，fast.ai架構(gòu)等等。

舉幾個(gè)例子好了。

比如說(shuō)，梯度下降。

原來(lái)，你需要這樣：

# clear last step
optimizer.zero_grad()
# 16 accumulated gradient steps
scaled_loss = 0
for accumulated_step_i in range(16):
 out = model.forward()
 loss = some_loss(out,y)
 loss.backward()
 scaled_loss += loss.item()
# update weights after 8 steps. effective batch = 8*16
optimizer.step()
# loss is now scaled up by the number of accumulated batches
actual_loss = scaled_loss / 16

在Lightning里，這一整段代碼不需要你自己敲了，只需輸入以下兩行代碼：

trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

不僅如此，在Lightning里，想用上單個(gè)GPU，直接調(diào)用即可：

trainer = Trainer(gpus = [0])
trainer.fit(model)

使用能將內(nèi)存占用減少一半的黑科技16位精度，不費(fèi)吹灰之力：

trainer = Trainer(amp_level = ’02’, use_amp = False)
trainer.fit(model)

使用多個(gè)GPU進(jìn)行分布式訓(xùn)練，so easy：

trainer = Trainer(gpus = [0, 1, 2, 3])
trainer.fit(model)

甚至是在1024個(gè)節(jié)點(diǎn)上以1024個(gè)GPU進(jìn)行訓(xùn)練，也是開箱即用：

trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7]）

自動(dòng)化功能還遠(yuǎn)不止這一些，以下模塊，均包含其中：

想要訓(xùn)練閃電那么快的神經(jīng)網(wǎng)絡(luò)嗎？Lightning簡(jiǎn)直為此量身定做。

此外，Lightning還和Tensorboard集成在了一起，可以輕松實(shí)現(xiàn)可視化學(xué)習(xí)。

只需記錄實(shí)驗(yàn)路徑：

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

然后在該路徑運(yùn)行tensorboard即可：

tensorboard —logdir /some/path

食用方法

想要使用Lightning，需要完成兩件事。

1、定義Lightning Model

這一步會(huì)花費(fèi)掉比較長(zhǎng)的時(shí)間。

Lightning Model是nn.Module的嚴(yán)格超類，它提供了與模型進(jìn)行交互的標(biāo)準(zhǔn)界面。

啟用Lightning Model最簡(jiǎn)單的方法是根據(jù)下面這個(gè)最小示例（minimal example）進(jìn)行局部修改:

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)

2、擬合訓(xùn)練器

訓(xùn)練器能處理Lightning自動(dòng)化部分的代碼核心邏輯，它會(huì)在訓(xùn)練過(guò)程中提取出最佳實(shí)踐。

基本的用法是像這樣：

from pytorch_lightning import Trainermoder = LightningTemplate()trainer = Trainer()
trainer.fit(model)

只要確保它的正確執(zhí)行，只需一個(gè)Trainer，計(jì)算集群（SLURM），Debug，分布式訓(xùn)練就通通不在話下了。

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')

One More Thing

你可能會(huì)問(wèn)，為什么要搞一個(gè)Lightning呢，用fast.ai不好嗎？

作者小哥表示，Lightning和fast.ai之間就沒(méi)什么好比的，fast.ai面向有志于進(jìn)入深度學(xué)習(xí)領(lǐng)域的新手，而Lightning面向的是ML領(lǐng)域中活躍的研究人員們。

就算真的要比，Lightning可是開箱即用的，不僅如此，在高性能計(jì)算、調(diào)試工具和可用性方面，小哥都對(duì)Lightning充滿信心。他自信地甩出了三張對(duì)比表格：

嗯，PyTorch真香！

傳送門

GitHub地址：

https://github.com/williamFalcon/pytorch-lightning

教學(xué)博客：

https://towardsdatascience.com/supercharge-your-ai-research-with-pytorch-lightning-337948a99eec

— 完 —