本文經(jīng)AI新媒體量子位（公眾號(hào)ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

對(duì)于PyTorch開(kāi)發(fā)者來(lái)說(shuō)，Google Colab上的TPU資源不能用，恐怕是最遺憾的事情了。Google的云計(jì)算資源眼看著不能用，轉(zhuǎn)TensorFlow又是不可能的。

過(guò)去一直有PyTorch用戶試圖在Colab上薅羊毛，但是都沒(méi)有太成功的。

現(xiàn)在福利來(lái)了，一個(gè)叫做Pytorch Lightning的項(xiàng)目，可以讓你幾乎修改代碼的情況下用上TPU。

Pytorch Lightning已經(jīng)上傳到PyPI，因此只需一行代碼就能安裝這個(gè)軟件。

pip install pytorch-lightning 

該項(xiàng)目的開(kāi)發(fā)者William Falcon說(shuō)，PyTorch Lightning是他在紐約大學(xué)和FAIR做博士生時(shí)研發(fā)，專門為從事AI研究的專業(yè)研究人員和博士生創(chuàng)建的。

幾乎無(wú)需修改代碼

首先讓我們來(lái)看一個(gè)MNIST圖像分類網(wǎng)絡(luò)的搭建，PyTorch的原始代碼和修改后的PyTorch Lightning代碼幾乎無(wú)異。

我們只需將nn.Module替換為pl.LightningModule即可。

作者表示，相比切換框架，用這種方法重構(gòu)原來(lái)的代碼只需數(shù)小時(shí)的時(shí)間。

PyTorch Lightning所做的是將科學(xué)代碼與工程代碼分離，只需將研究代碼重構(gòu)為L(zhǎng)ightningModule格式（科學(xué)），Lightning將自動(dòng)完成其余部分（工程）。

使用方法

PyTorch Lightning具體該如何使用，作者Falcon還是以MNIST圖像分類網(wǎng)絡(luò)為例，介紹從收集數(shù)據(jù)到訓(xùn)練再到驗(yàn)證、測(cè)試的全過(guò)程。

準(zhǔn)備數(shù)據(jù)集階段分為下載圖片、轉(zhuǎn)換、分割數(shù)據(jù)集、打包四個(gè)步驟。二者代碼大致相同，只是將PyTorch代碼組織為4個(gè)函數(shù)：

prepare_data：此函數(shù)負(fù)責(zé)處理下載數(shù)據(jù)，確保使用多個(gè)GPU時(shí)，不會(huì)下載多個(gè)數(shù)據(jù)集或?qū)?shù)據(jù)進(jìn)行雙重操作。

train_dataloader，val_dataloader，test_dataloader：每一個(gè)都負(fù)責(zé)返回相應(yīng)數(shù)據(jù)集的數(shù)據(jù)拆分。以這種方式進(jìn)行構(gòu)造訓(xùn)練、驗(yàn)證、測(cè)試集，可以讓你非常清楚如何操作數(shù)據(jù)。

接下來(lái)是優(yōu)化器的選擇，比如選擇Adam。兩者的代碼還是幾乎完全相同，不過(guò)后者把pytorch_model改成了self。

至于損失函數(shù)，對(duì)于n向分類，要使用交叉熵?fù)p失。二者代碼又幾乎一致相同，后者多出一個(gè)self。

在訓(xùn)練上，PyTorch Lightning的代碼更簡(jiǎn)潔一點(diǎn)。在PyTorch中，你需要自己編寫(xiě)for循環(huán)，這意味著你必須記住要以正確的順序調(diào)用正確的東西，可能會(huì)導(dǎo)致錯(cuò)誤。

而PyTorch Lightning將樣板格式抽象化，但方框中的內(nèi)容保持不變，增加了代碼的可讀性和可重復(fù)性。

后面驗(yàn)證推理的部分不再贅述。

總之，PyTorch Lightning有這些優(yōu)點(diǎn)：

代碼結(jié)構(gòu)化；與PyTorch源代碼幾乎完全相同；隨著項(xiàng)目復(fù)雜性的提升，代碼的大部分內(nèi)容無(wú)需修改；保留了PyTorch的靈活性。

新增高級(jí)功能，連Pytorch本尊都沒(méi)有

除了以上的一些特性外，PyTorch Lightning還加入了許多高級(jí)功能，讓你體驗(yàn)到PyTorch本身不具備的一些優(yōu)點(diǎn)。

比如更清晰直觀的訓(xùn)練進(jìn)度條：

用TensorBoard日志記錄代碼運(yùn)行全過(guò)程：

PyTorch Lightning還支持TensorBoard之外的其它5種工具記錄日志：

與TensorBoard，MLFlow完全集成，并支持任何日志記錄模塊。

甚至還有一個(gè)內(nèi)置的分析工具，告訴你訓(xùn)練過(guò)程中的瓶頸：

trainer = Trainer(…, profiler=True) 

PyTorch Lightning還有更多的可擴(kuò)展性，在這里無(wú)法一一介紹，如果你正想要在TPU上運(yùn)行自己的PyTorch代碼，可以前去學(xué)習(xí)更詳細(xì)的用法。

傳送門

項(xiàng)目地址：

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

Colab演示：

https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3#scrollTo=dEeUzX_5aLrX