如何提升 PyTorch「煉丹」速度?
最近,知名機器學習與 AI 研究者 Sebastian Raschka 向我們展示了他的絕招。據他表示,他的方法在不影響模型準確率的情況下,僅僅透過改變幾行程式碼,將 BERT 優化時間從 22.63 分鐘縮減到 3.15 分鐘,訓練速度足足提升了 7 倍。
#作者更是表示,如果你有8 個GPU 可用,整個訓練過程只需要2 分鐘,實現11.5 倍的效能加速。
#下面我們來看看他到底是如何實現的。
首先是模型,作者採用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中文網其他相關文章!