改變幾行程式碼，PyTorch煉丹速度狂飆、模型最佳化時間大減

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

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

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

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

讓PyTorch 模型訓練更快

首先是模型，作者採用DistilBERT 模型進行研究，它是BERT 的精簡版，與BERT 相比規模縮小了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 模型訓練速度的方法，有興趣的夥伴可以跟著原部落格嘗試一下，相信你會得到想要的結果。 ######

以上是改變幾行程式碼，PyTorch煉丹速度狂飆、模型最佳化時間大減的詳細內容。更多資訊請關注PHP中文網其他相關文章！