基於PyTorch的「Keras」：除了核心邏輯通通都封裝

機器之心報道

參與：思源、一鳴

Keras 和 PyTorch 都是對初學者最友好的深度學習框架，它們用起來就像描述架構的簡單語言一樣，告訴框架哪一層該用什麼就行。很多研究者和開發者都在考慮到底哪一個框架更好，但目前兩個框架都非常流行，它們都各有優勢。最近，Facebook 研究員 William Falcon 為 PyTorch 披上了一件 Keras 的外衣，他表明用這樣的框架做研究簡直不要太爽。

PyTorch Lightning 地址：https://github.com/williamFalcon/pytorch-lightning

看起來像 Keras 的 PyTorch

Keras 本身的目的就是對深度學習框架（TensorFlow、Theano）進行了進一步的 API 封裝。作為 TensorFlow 的高度封裝，Keras 的抽象層次非常高，很多 API 細節都隱藏了起來。雖然 PyTorch 同樣使用動態計算圖，也方便快捷，但總體上 Keras 隱藏的細節更多一些。

反觀 PyTorch，它提供一個相對較低級別的實驗環境，使用戶可以更加自由地編寫自定義層、查看數值優化任務等等。例如在 PyTorch 1.0 中，編譯工具 torch.jit 就包含一種名為 Torch Script 的語言，它是 Python 的子語言，開發者使用它能進一步對模型進行優化。

用 PyTorch 寫模型，除了數據載入和模型定義部分外，整個訓練和驗證的邏輯、配置都需要我們手動完成，這些步驟都較為繁瑣。甚至可以說，研究者需要耗費相當多的精力處理這一部分的代碼，還要祈禱不出 Bug。但是對於大多數研究實驗來說，訓練和驗證的循環體差不多都是一樣的，實現的功能也相當一致，所以為什麼不將這些通用的東西都打包在一起，這樣訓練不就簡單了么？

William Falcon 正是這樣想的，他將 PyTorch 開發中的各種通用配置全都包裝起來，我們只需要寫核心邏輯就行。通過 PyTorch Lightning，PyTorch 就類似於 Keras，它能以更高級的形式快速搭建模型。

項目作者是誰

要完成這樣的工作，工作量肯定是非常大的，因為從超參搜索、模型 Debug、分散式訓練、訓練和驗證的循環邏輯到模型日誌的列印，都需要寫一套通用的方案，確保各種任務都能用得上。所以 Facebook 的這位小哥哥 William Falcon 還是很厲害的。

他是一位 NYU 和 Facebook 的開發者。目前在 NYU 攻讀 PhD。從 GitHub 的活動來看，小哥是一位比較活躍的開發者。

這是一位披著 Keras 外衣的 PyTorch

Lightning 是 PyTorch 非常輕量級的包裝，研究者只需要寫最核心的訓練和驗證邏輯，其它過程都會自動完成。因此這就有點類似 Keras 那種高級包裝，它隱藏了絕大多數細節，只保留了最通俗易懂的介面。Lightning 能確保自動完成部分的正確性，對於核心訓練邏輯的提煉非常有優勢。

那麼我們為什麼要用 Lightning？

當我們開始構建新項目，最後你希望做的可能就是記錄訓練循環、多集群訓練、float16 精度、提前終止、模型載入/保存等等。這一系列過程可能需要花很多精力來解決各式各樣、千奇百怪的 Bug，因此很難把精力都放在研究的核心邏輯上。

通過使用 Lightning，這些部分都能保證是 Work 的，因此能抽出精力關注我們要研究的東西：數據、訓練、驗證邏輯。此外，我們完全不需要擔心使用多 GPU 加速會很難，因為 Lightning 會把這些東西都做好。

所以 Lightning 都能幫我們幹什麼？

下圖展示了構建一個機器學習模型都會經歷哪些過程，很多時候最困難的還不是寫模型，是各種配置與預處理過程。如下藍色的部分需要用 LightningModule 定義，而灰色部分 Lightning 可以自動完成。我們需要做的，差不多也就載入數據、定義模型、確定訓練和驗證過程。

下面的偽代碼展示了大致需要定義的幾大模塊，它們再加上模型架構定義就能成為完整的模型。

除了需要定義的模塊外，以下步驟均可通過 Lightning 自動完成。當然，每個模塊可以單獨進行配置。

Lightning 怎麼用

Lightning 的使用也非常簡單，只需要兩步就能完成：定義 LightningModel；擬合訓練器。

以經典的 MNIST 圖像識別為例，如下展示了 LightningModel 的示例。我們可以照常導入 PyTorch 模塊，但這次不是繼承 nn.Module，而是繼承 LightningModel。然後我們只需要照常寫 PyTorch 就行了，該調用函數還是繼續調用。這裡看上去似乎沒什麼不同，但注意方法名都是確定的，這樣才能利用 Lightning 的後續過程。

隨後，第二步即擬合訓練器。這就比較類似 Keras 這類高級包裝了，它將訓練配置細節、循環體、以及日誌輸出等更加具體的信息全都隱藏了，一個 fit() 方法就能自動搞定一切。

這相比以前寫 PyTorch 更加便捷精鍊一些，而且分散式訓練也非常容易，只要給出設備 id 就行了。

其他特性

Pytorch-Lightning 還可以和 TensorBoard 無縫對接。

只需要定義運行的路徑：

將 TensorBoard 連接到路徑上即可：

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！