輕量化模型訓練加速的思考（Pytorch實現）

LyleChen發表於2020-09-01
模型PyTorch

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 伺服器

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

2.2 基本配置

  • Dataset：ImageNet
  • Model：MobilenetV2
  • Augmentation：RandomCrop，RandomFlip，Resize，Normalization
  • 每個程式的batch_size：256
  • 每個程式的Dataloader的num_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_gpusnum_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. 一些總結

通過綜合採用各種訓練加速手段，基本可以做到充分利用多顯示卡伺服器的GPU和CPU資源，不會造成硬體資源的浪費；

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

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

6. 一些意外

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

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

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

下一步是要折騰水冷散熱了麼？

相關文章