## 0. 引子 在訓練輕量化模型時，經常發生的情況就是，明明 GPU 很閒，可速度就是上不去，用了多張卡並行也沒有太大改善。 如果什麼優化都不做，僅僅是使用`nn.DataParallel`這個模組，那麼實測大概只能實現一點幾倍的加速（按每秒處理的總圖片數計算），不管用多少張卡。因為卡越多，資料傳輸的開銷就越大，副作用就越大。 為了提高GPU伺服器的資源利用率，嘗試了一些加速的手段。 > 基於Pytorch1.6.0版本實現，官方支援`amp`功能，不再需要外部`apex`庫； > 此外比較重要的庫是Dali。 梳理了訓練框架，並將參考程式碼放到Github上。 如果覺得對你有所啟發，請給個star呀。 [**參考程式碼**](https://github.com/LyleC/PytorchTrainAcceleration) ## 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_gpus`*`num_threads` < CPU核心數，不能充分利用CPU資源； 如果`num_gpus`*`num_threads` > CPU核心數，速度反而也會有所下降。 ## 4. 訓練速度實測結果 ### 4.1 未開啟amp時的GPU佔用  ### 4.2 開啟amp後的GPU佔用  ### 4.3 CPU佔用情況 > 開啟/關閉amp對於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利用率直接降到了原來的一半左右。 之前顯示卡執行效率低的時候，散熱不良的問題還沒有顯露出來，一旦長時間高負荷運轉，多卡密集排布和風冷散熱的不足就暴露出來了。 下一步是要折騰水冷散熱