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

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

過去一直有PyTorch用戶試圖在Colab上薅羊毛，但是都沒有太成功的。

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

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

pip install pytorch-lightning 

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

幾乎無需修改代碼

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

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

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

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

使用方法

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

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

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

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

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

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

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

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

后面驗證推理的部分不再贅述。

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

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

新增高級功能，連Pytorch本尊都沒有

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

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

用TensorBoard日志記錄代碼運行全過程：

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

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

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

trainer = Trainer(…, profiler=True) 

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

傳送門

項目地址：

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

Colab演示：

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