讓PyTorch 使用更輕便的深度學(xué)習(xí)框架

一直以來，PyTorch就以簡(jiǎn)單又好用的特點(diǎn)，廣受AI研究者的喜愛。 但是，一旦任務(wù)復(fù)雜化，就可能會(huì)發(fā)生一系列錯(cuò)誤，花費(fèi)的時(shí)間更長(zhǎng)。 于是，就誕生了這樣一個(gè)“友好”的PyTorch Lightning。

直接在GitHub上斬獲6.6k星。

首先，它把研究代碼與工程代碼相分離，還將PyTorch代碼結(jié)構(gòu)化，更加直觀的展現(xiàn)數(shù)據(jù)操作過程。 這樣，更加易于理解，不易出錯(cuò)，本來很冗長(zhǎng)的代碼一下子就變得輕便了，對(duì)AI研究者十分的友好。 話不多說，我們就來看看這個(gè)輕量版的“PyTorch”。

關(guān)于Lightning

Lightning將DL/ML代碼分為三種類型：研究代碼、工程代碼、非必要代碼。 針對(duì)不同的代碼，Lightning有不同的處理方式。 這里的研究代碼指的是特定系統(tǒng)及其訓(xùn)練方式，比如GAN、VAE，這類的代碼將由LightningModule直接抽象出來。 我們以MNIST生成為例。

l1 = nn.Linear(...)l2 = nn.Linear(...)decoder = Decoder()
x1 = l1(x)x2 = l2(x2)out = decoder(features, x)
loss = perceptual_loss(x1, x2, x) + CE(out, x)

而工程代碼是與培訓(xùn)此系統(tǒng)相關(guān)的所有代碼，比如提前停止、通過GPU分配、16位精度等。 我們知道，這些代碼在大多數(shù)項(xiàng)目中都相同，所以在這里，直接由Trainer抽象出來。

model.cuda(0)x = x.cuda(0)
distributed = DistributedParallel(model)
with gpu_zero:download_data()
dist.barrier()

剩下的就是非必要代碼，有助于研究項(xiàng)目，但是與研究項(xiàng)目無關(guān)，可能是檢查梯度、記錄到張量板。此代碼由Callbacks抽象出來。

# log samplesz = Q.rsample()generated = decoder(z)self.experiment.log('images', generated)

此外，它還有一些的附加功能，比如你可以在CPU，GPU，多個(gè)GPU或TPU上訓(xùn)練模型，而無需更改PyTorch代碼的一行；你可以進(jìn)行16位精度訓(xùn)練，可以使用Tensorboard的五種方式進(jìn)行記錄。 這樣說，可能不太明顯，我們就來直觀的比較一下PyTorch與PyTorch Lightning之間的差別吧。

PyTorch與PyTorch Lightning比較

直接上圖。

我們就以構(gòu)建一個(gè)簡(jiǎn)單的MNIST分類器為例，從模型、數(shù)據(jù)、損失函數(shù)、優(yōu)化這四個(gè)關(guān)鍵部分入手。

模型

首先是構(gòu)建模型，本次設(shè)計(jì)一個(gè)3層全連接神經(jīng)網(wǎng)絡(luò)，以28×28的圖像作為輸入，將其轉(zhuǎn)換為數(shù)字0-9的10類的概率分布。

兩者的代碼完全相同。意味著，若是要將PyTorch模型轉(zhuǎn)換為PyTorch Lightning，我們只需將nn.Module替換為pl.LightningModule。 也許這時(shí)候，你還看不出這個(gè)Lightning的神奇之處。不著急，我們接著看。

數(shù)據(jù)

接下來是數(shù)據(jù)的準(zhǔn)備部分，代碼也是完全相同的，只不過Lightning做了這樣的處理。 它將PyTorch代碼組織成了4個(gè)函數(shù)，prepare_data、train_dataloader、val_dataloader、test_dataloader。

prepare_data 這個(gè)功能可以確保在你使用多個(gè)GPU的時(shí)候，不會(huì)下載多個(gè)數(shù)據(jù)集或者對(duì)數(shù)據(jù)進(jìn)行多重操作。這樣所有代碼都確保關(guān)鍵部分只從一個(gè)GPU調(diào)用。 這樣就解決了PyTorch老是重復(fù)處理數(shù)據(jù)的問題，這樣速度也就提上來了。 train_dataloader, val_dataloader, test_dataloader 每一個(gè)都負(fù)責(zé)返回相應(yīng)的數(shù)據(jù)分割，這樣就能很清楚的知道數(shù)據(jù)是如何被操作的，在以往的教程里，都幾乎看不到它們的是如何操作數(shù)據(jù)的。 此外，Lightning還允許使用多個(gè)dataloaders來測(cè)試或驗(yàn)證。

優(yōu)化

接著就是優(yōu)化。

不同的是，Lightning被組織到配置優(yōu)化器的功能中。如果你想要使用多個(gè)優(yōu)化器，則可同時(shí)返回兩者。

損失函數(shù)

對(duì)于n向分類，我們要計(jì)算交叉熵?fù)p失。兩者的代碼是完全一樣的。

此外，還有更為直觀的——驗(yàn)證和訓(xùn)練循環(huán)。

在PyTorch中，我們知道，需要你自己去構(gòu)建for循環(huán)，可能簡(jiǎn)單的項(xiàng)目還好，但是一遇到更加復(fù)雜高級(jí)的項(xiàng)目就很容易翻車了。 而Lightning里這些抽象化的代碼，其背后就是由Lightning里強(qiáng)大的trainer團(tuán)隊(duì)負(fù)責(zé)了。

PyTorch Lightning安裝教程

看到這里，是不是也想安裝下來試一試。 PyTorch Lightning安裝十分簡(jiǎn)單。 代碼如下：

conda activate my_envpip install pytorch-lightning

或在沒有conda環(huán)境的情況下，可以在任何地方使用pip。 代碼如下：

pip install pytorch-lightning

創(chuàng)建者也有大來頭

William Falcon，PyTorch Lightning 的創(chuàng)建者，現(xiàn)在在紐約大學(xué)的人工智能專業(yè)攻讀博士學(xué)位，在《福布斯》擔(dān)任AI特約作者。 2018年，從哥倫比亞大學(xué)計(jì)算機(jī)科學(xué)與統(tǒng)計(jì)學(xué)專業(yè)畢業(yè)，本科期間，他還曾輔修數(shù)學(xué)。 現(xiàn)在已獲得Google Deepmind資助攻讀博士學(xué)位的獎(jiǎng)學(xué)金，去年還收到Facebook AI Research實(shí)習(xí)邀請(qǐng)。 此外，他還曾是一個(gè)海軍軍官，接受過美國(guó)海軍海豹突擊隊(duì)的訓(xùn)練。

前不久，華爾街日?qǐng)?bào)就曾還曾提到這個(gè)團(tuán)隊(duì)，他們正在研究呼吸系統(tǒng)疾病與呼吸模式之間的聯(lián)系?？赡軙?huì)應(yīng)用到的場(chǎng)景，是通過電話在診斷新冠癥狀。目前，該團(tuán)隊(duì)還處在數(shù)據(jù)收集階段。 果然，優(yōu)秀的人，干什么都是優(yōu)秀的。嘆氣……