本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

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

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

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

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

pip install pytorch-lightning 

該項目的開發者William Falcon說，PyTorch Lightning是他在紐約大學和FAIR做博士生時研發，專門為從事AI研究的專業研究人員和博士生創建的。

幾乎無需修改代碼

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

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

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

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

使用方法

PyTorch Lightning具體該如何使用，作者Falcon還是以MNIST圖像分類網絡為例，介紹從收集數據到訓練再到驗證、測試的全過程。

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

prepare_data：此函數負責處理下載數據，確保使用多個GPU時，不會下載多個數據集或對數據進行雙重操作。

train_dataloader，val_dataloader，test_dataloader：每一個都負責返回相應數據集的數據拆分。以這種方式進行構造訓練、驗證、測試集，可以讓你非常清楚如何操作數據。

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

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

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

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

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

總之，PyTorch Lightning有這些優點：

代碼結構化；與PyTorch源代碼幾乎完全相同；隨著項目復雜性的提升，代碼的大部分內容無需修改；保留了PyTorch的靈活性。

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

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

比如更清晰直觀的訓練進度條：

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

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

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

甚至還有一個內置的分析工具，告訴你訓練過程中的瓶頸：

trainer = Trainer(…, profiler=True) 

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

傳送門

項目地址：

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

Colab演示：

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