改變幾行代碼，PyTorch煉丹速度狂飆、模型優(yōu)化時間大減
如何提升 PyTorch「煉丹」速度？
最近，知名機器學習與 AI 研究者 Sebastian Raschka 向我們展示了他的絕招。據(jù)他表示，他的方法在不影響模型準確率的情況下，僅僅通過改變幾行代碼，將 BERT 優(yōu)化時間從 22.63 分鐘縮減到 3.15 分鐘，訓練速度足足提升了 7 倍。
作者更是表示，如果你有 8 個 GPU 可用，整個訓練過程只需要 2 分鐘，實現(xiàn) 11.5 倍的性能加速。
下面我們來看看他到底是如何實現(xiàn)的。
讓 PyTorch 模型訓練更快
首先是模型，作者采用 DistilBERT 模型進行研究，它是 BERT 的精簡版，與 BERT 相比規(guī)模縮小了 40%，但性能幾乎沒有損失。其次是數(shù)據(jù)集，訓練數(shù)據(jù)集為大型電影評論數(shù)據(jù)集 IMDB Large Movie Review，該數(shù)據(jù)集總共包含 50000 條電影評論。作者將使用下圖中的 c 方法來預測數(shù)據(jù)集中的影評情緒。
基本任務交代清楚后，下面就是 PyTorch 的訓練過程。為了讓大家更好地理解這項任務，作者還貼心地介紹了一下熱身練習，即如何在 IMDB 電影評論數(shù)據(jù)集上訓練 DistilBERT 模型。如果你想自己運行代碼，可以使用相關的 Python 庫設置一個虛擬環(huán)境，如下所示：
相關軟件的版本如下：
現(xiàn)在省略掉枯燥的數(shù)據(jù)加載介紹，只需要了解本文將數(shù)據(jù)集劃分為 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，它定義了如何執(zhí)行訓練、驗證和測試。相比于前面給出的代碼，主要變化是在第 5 部分（即 ### 5 Finetuning），即微調模型。與以前不同的是，微調部分在 LightningModel 類中包裝了 PyTorch 模型，并使用 Trainer 類來擬合模型。
之前的代碼顯示驗證準確率從第 2 輪到第 3 輪有所下降，但改進后的代碼使用了 ModelCheckpoint 以加載最佳模型。在同一臺機器上，這個模型在 23.09 分鐘內達到了 92% 的測試準確率。
需要注意，如果禁用 checkpointing 并允許 PyTorch 以非確定性模式運行，本次運行最終將獲得與普通 PyTorch 相同的運行時間（時間為 22.63 分而不是 23.09 分）。
自動混合精度訓練
進一步，如果 GPU 支持混合精度訓練，可以開啟 GPU 以提高計算效率。作者使用自動混合精度訓練，在 32 位和 16 位浮點之間切換而不會犧牲準確率。
在這一優(yōu)化下，使用 Trainer 類，即能通過一行代碼實現(xiàn)自動混合精度訓練：
上述操作可以將訓練時間從 23.09 分鐘縮短到 8.75 分鐘，這幾乎快了 3 倍。測試集的準確率為 92.2%，甚至比之前的 92.0% 還略有提高。
使用 Torch.Compile 靜態(tài)圖
最近 PyTorch 2.0 公告顯示，PyTorch 團隊引入了新的 toch.compile 函數(shù)。該函數(shù)可以通過生成優(yōu)化的靜態(tài)圖來加速 PyTorch 代碼執(zhí)行，而不是使用動態(tài)圖運行 PyTorch 代碼。
由于 PyTorch 2.0 尚未正式發(fā)布，因而必須先要安裝 torchtriton，并更新到 PyTorch 最新版本才能使用此功能。
然后通過添加這一行對代碼進行修改：
在 4 塊 GPU 上進行分布式數(shù)據(jù)并行
上文介紹了在單 GPU 上加速代碼的混合精度訓練，接下來介紹多 GPU 訓練策略。下圖總結了幾種不同的多 GPU 訓練技術。
想要實現(xiàn)分布式數(shù)據(jù)并行，可以通過 DistributedDataParallel 來實現(xiàn)，只需修改一行代碼就能使用 Trainer。
經過這一步優(yōu)化，在 4 個 A100 GPU 上，這段代碼運行了 3.52 分鐘就達到了 93.1% 的測試準確率。
DeepSpeed
最后，作者探索了在 Trainer 中使用深度學習優(yōu)化庫 DeepSpeed 以及多 GPU 策略的結果。首先必須安裝 DeepSpeed 庫：
接著只需更改一行代碼即可啟用該庫：
這一波下來，用時 3.15 分鐘就達到了 92.6% 的測試準確率。不過 PyTorch 也有 DeepSpeed 的替代方案：fully-sharded DataParallel，通過 strategy="fsdp" 調用，最后花費 3.62 分鐘完成。
以上就是作者提高 PyTorch 模型訓練速度的方法，感興趣的小伙伴可以跟著原博客嘗試一下，相信你會得到想要的結果。
