一行代碼安裝，TPU也能運行PyTorch，修改少量代碼即可快速移植
曉查 發自 凹非寺
量子位 報道 | 公眾號 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
- 腦機接口走向現實，11張PPT看懂中國腦機接口產業現狀｜量子位智庫2021-08-10
- 張朝陽開課手推E=mc2，李永樂現場狂做筆記2022-03-11
- 阿里數學競賽可以報名了！獎金增加到400萬元，題目面向大眾公開征集2022-03-14
- 英偉達遭黑客最后通牒：今天必須開源GPU驅動，否則公布1TB機密數據2022-03-05
相關閱讀
Kaldi拜拜！PyTorch語音工具包SpeechBrain要來了，支持多種語音任務，實現最強水準
能用來做語音識別、說話人識別、語音分離，多麥克風信號處理、自我監督和無監督學習、語音增強等
