如何提升 PyTorch「煉丹」速度？

最近，知名機器學習與 AI 研究者 Sebastian Raschka 向我們展示了他的絕招。據他表示，他的方法在不影響模型準確率的情況下，僅僅通過改變幾行代碼，將 BERT 優化時間從 22.63 分鐘縮減到 3.15 分鐘，訓練速度足足提升了 7 倍。

作者更是表示，如果你有 8 個 GPU 可用，整個訓練過程只需要 2 分鐘，實現 11.5 倍的性能加速。

下面我們來看看他到底是如何實現的。

讓 PyTorch 模型訓練更快

首先是模型，作者采用 DistilBERT 模型進行研究，它是 BERT 的精簡版，與 BERT 相比規?？s小了 40%，但性能幾乎沒有損失。其次是數據集，訓練數據集為大型電影評論數據集 IMDB Large Movie Review，該數據集總共包含 50000 條電影評論。作者將使用下圖中的 c 方法來預測數據集中的影評情緒。

基本任務交代清楚后，下面就是 PyTorch 的訓練過程。為了讓大家更好地理解這項任務，作者還貼心地介紹了一下熱身練習，即如何在 IMDB 電影評論數據集上訓練 DistilBERT 模型。如果你想自己運行代碼，可以使用相關的 Python 庫設置一個虛擬環境，如下所示：

相關軟件的版本如下：

現在省略掉枯燥的數據加載介紹，只需要了解本文將數據集劃分為 35000 個訓練示例、5000 個驗證示例和 10000 個測試示例。需要的代碼如下：

代碼部分截圖

完整代碼地址：

?https://github.com/rasbt/faster-pytorch-blog/blob/main/1_pytorch-distilbert.py?

然后在 A100 GPU 上運行代碼，得到如下結果：

部分結果截圖

正如上述代碼所示，模型從第 2 輪到第 3 輪開始有一點過擬合，驗證準確率從 92.89% 下降到了 92.09%。在模型運行了 22.63 分鐘后進行微調，最終的測試準確率為 91.43%。

使用 Trainer 類?

接下來是改進上述代碼，改進部分主要是把 PyTorch 模型包裝在 LightningModule 中，這樣就可以使用來自 Lightning 的 Trainer 類。部分代碼截圖如下：

完整代碼地址：https://github.com/rasbt/faster-pytorch-blog/blob/main/2_pytorch-with-trainer.py?

上述代碼建立了一個 LightningModule，它定義了如何執行訓練、驗證和測試。相比于前面給出的代碼，主要變化是在第 5 部分（即 ### 5 Finetuning），即微調模型。與以前不同的是，微調部分在 LightningModel 類中包裝了 PyTorch 模型，并使用 Trainer 類來擬合模型。

之前的代碼顯示驗證準確率從第 2 輪到第 3 輪有所下降，但改進后的代碼使用了 ModelCheckpoint 以加載最佳模型。在同一臺機器上，這個模型在 23.09 分鐘內達到了 92% 的測試準確率。

需要注意，如果禁用 checkpointing 并允許 PyTorch 以非確定性模式運行，本次運行最終將獲得與普通 PyTorch 相同的運行時間（時間為 22.63 分而不是 23.09 分）。

自動混合精度訓練

進一步，如果 GPU 支持混合精度訓練，可以開啟 GPU 以提高計算效率。作者使用自動混合精度訓練，在 32 位和 16 位浮點之間切換而不會犧牲準確率。

在這一優化下，使用 Trainer 類，即能通過一行代碼實現自動混合精度訓練：

上述操作可以將訓練時間從 23.09 分鐘縮短到 8.75 分鐘，這幾乎快了 3 倍。測試集的準確率為 92.2%，甚至比之前的 92.0% 還略有提高。

使用 Torch.Compile 靜態圖

最近 PyTorch 2.0 公告顯示，PyTorch 團隊引入了新的 toch.compile 函數。該函數可以通過生成優化的靜態圖來加速 PyTorch 代碼執行，而不是使用動態圖運行 PyTorch 代碼。

由于 PyTorch 2.0 尚未正式發布，因而必須先要安裝 torchtriton，并更新到 PyTorch 最新版本才能使用此功能。

然后通過添加這一行對代碼進行修改：

在 4 塊 GPU 上進行分布式數據并行

上文介紹了在單 GPU 上加速代碼的混合精度訓練，接下來介紹多 GPU 訓練策略。下圖總結了幾種不同的多 GPU 訓練技術。

想要實現分布式數據并行，可以通過 DistributedDataParallel 來實現，只需修改一行代碼就能使用 Trainer。

經過這一步優化，在 4 個 A100 GPU 上，這段代碼運行了 3.52 分鐘就達到了 93.1% 的測試準確率。

DeepSpeed

最后，作者探索了在 Trainer 中使用深度學習優化庫 DeepSpeed 以及多 GPU 策略的結果。首先必須安裝 DeepSpeed 庫：

接著只需更改一行代碼即可啟用該庫：

這一波下來，用時 3.15 分鐘就達到了 92.6% 的測試準確率。不過 PyTorch 也有 DeepSpeed 的替代方案：fully-sharded DataParallel，通過 strategy="fsdp" 調用，最后花費 3.62 分鐘完成。

以上就是作者提高 PyTorch 模型訓練速度的方法，感興趣的小伙伴可以跟著原博客嘗試一下，相信你會得到想要的結果。