jlzzjlzz亚洲乱熟在线播放

系統城裝機大師 - 唯一官網:www.farandoo.com!

當前位置:首頁 > 腳本中心 > python > 詳細頁面

輕量化模型訓練加速的思考(Pytorch實現)

時間:2020-09-01來源:www.farandoo.com作者:電腦系統城

0. 引子

在訓練輕量化模型時,經常發生的情況就是,明明 GPU 很閑,可速度就是上不去,用了多張卡并行也沒有太大改善。

如果什么優化都不做,僅僅是使用nn.DataParallel這個模塊,那么實測大概只能實現一點幾倍的加速(按每秒處理的總圖片數計算),不管用多少張卡。因為卡越多,數據傳輸的開銷就越大,副作用就越大。

為了提高GPU服務器的資源利用率,嘗試了一些加速的手段。

基于Pytorch1.6.0版本實現,官方支持amp功能,不再需要外部apex庫;
此外比較重要的庫是Dali。

梳理了訓練框架,并將參考代碼放到Github上。
如果覺得對你有所啟發,請給個star呀。

參考代碼

1. 訓練速度的瓶頸及應對思路

這邊主要說的是CV領域,但在其他領域,思路應該也是相通的。
模型訓練過程中,影響整體速度的因素主要有以下幾點:

  1. 將數據從磁盤讀取到內存的效率;
  2. 對圖片進行解碼的效率;
  3. 對樣本進行在線增強的效率;
  4. 網絡前向/反向傳播和Loss計算的效率;

針對這幾個因素,分別采取如下幾種應對思路:

  1. 加快數據讀取可以有幾種思路:
    • 采取類似TF的tfrecord或者Caffe的lmdb格式,提前將數據打包,比反復加載海量的小文件要快很多,但pytorch沒有通用的數據打包方式;
    • 在初始化時,提前將所有數據加載到內存中(前提是數據集不能太大,內存能裝得下);
    • 將數據放在SSD而非HDD,可以大大提速(前提是你有足夠大的SSD);
  2. 提升圖片解碼速度,可以考慮采用NVIDIA-DALI庫,能夠利用GPU來加速JPG格式的圖片解碼,針對其他格式的圖片(如PNG),不能實現GPU加速,但也可以兼容;
  3. 提升樣本在線增強的效率,同樣可以通過NVIDIA-DALI庫,實現GPU加速;
  4. 在網絡結構確定的情況下,提速主要有兩種方式,并且可以同時采用:
    • 采用Data Parallel的多卡并行訓練
    • 采用amp自動混合精度訓練

2. 實驗配置

2.1 服務器

服務器為4TITAN RTX,進行實驗時停止了其他高資源消耗的進程。

2.2 基本配置

  • Dataset:ImageNet
  • Model:MobilenetV2
  • Augmentation:RandomCrop,RandomFlip,Resize,Normalization
  • 每個進程的batch_size:256
  • 每個進程的Dataloadernum_threads:8

3. 具體實現中的注意點

3.1 關于Dataloader

在使用DALI庫構建Dataloader時,建議采用ops.ExternalSource的方式來加載數據,這樣可以實現比較高的自由度。

示例代碼中只實現了分類任務的dataloader,但按照這種方式構建,很容易實現其他如檢測/分割任務的dataloader。

只要把數據來源按照迭代器來實現,就可以套用到ops.ExternalSource這一套框架下。

參見src/datasets/cls_dataset_dali.py中的ClsInputIterator。

3.2 關于Loss

在訓練過程中,每個進程分別計算各自的loss,通過內部同步機制去同步loss信息。但是在訓練中需要監控過程,此時需要計算所有loss的均值。

參見src/train/logger.py中關于reduce_tensor的計算方式。

3.3 關于多進程參數的選取

在訓練過程中,實驗用的服務器,CPU共32核心,4卡并行,因此每個進程的Dataloader,設定的num_threads為8,測試下來效率最高。
如果num_gpusnum_threads < CPU核心數,不能充分利用CPU資源;
如果num_gpus
num_threads > CPU核心數,速度反而也會有所下降。

4. 訓練速度實測結果

4.1 未開啟amp時的GPU占用

未開啟amp功能

4.2 開啟amp后的GPU占用

開啟amp功能

4.3 CPU占用情況

開啟/關閉amp對于CPU的影響不大,基本看不出區別

CPU

4.4 綜合訓練速度

4卡并行,BS為256,訓練集約120W圖片。訓練速度為:

  • 未開啟amp:約 2.4 iters/s(2458 幀/s),每個epoch訓練時間不到 9 min;
  • 開啟amp:約 3.8 iters/s(3891 幀/s),每個epoch訓練時間不到 6 min;

5. 一些總結

通過綜合采用各種訓練加速手段,基本可以做到充分利用多顯卡服務器的GPUCPU資源,不會造成硬件資源的浪費;

  1. 通過Nvidia-Dali模塊的合理配置,可以顯著提升數據加載和在線增強階段的效率,特別是在訓練一些輕量化模型時,往往瓶頸不在于GPU的計算速度,而在于CPU等其他部件的負載;
  2. 通過DistributedDataParallel模塊的合理配置,可以實現多卡的負載均衡,不論是顯存占用還是GPU利用率,都能夠達到平衡,不會有其中1張卡變成效率瓶頸;
  3. 通過torch.cuda.amp模塊的合理配置,可以進一步降低顯存占用,從而可以設置更大的batch_size,提高模型收斂速度;
  4. torch.cuda.amp模塊還可以顯著降低網絡前向推理時間,從而進一步提高訓練效率。

綜合應用如上所述的手段,基本上可以實現顯卡數量和訓練效率之間的線性增長關系。
不會發生卡多了,但是單卡的效率卻大大下降的現象。

6. 一些意外

原以為本篇到此就該結束了,但又遇到了新的問題。

當訓練執行一段時間后,由于整個系統長時間處于高負載的狀態,顯卡溫度飆升,觸發了顯卡的保護機制,自動降頻了,GPU利用率直接降到了原來的一半左右。

之前顯卡運行效率低的時候,散熱不良的問題還沒有顯露出來,一旦長時間高負荷運轉,多卡密集排布和風冷散熱的不足就暴露出來了。

下一步是要折騰水冷散熱了么?

分享到:

相關信息

系統教程欄目

欄目熱門教程

人氣教程排行

站長推薦

熱門系統下載