曉查 發自 凹非寺量子位 報導 | 公眾號 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