0. 引子
在訓練輕量化模型時，經常發生的情況就是，明明 GPU 很閑，可速度就是上不去，用了多張卡並行也沒有太大改善。
如果什么優化都不做，僅僅是使用nn.DataParallel這個模塊，那么實測大概只能實現一點幾倍的加速（按每秒處理的總圖片數計算），不管用多少張卡。因為卡越多，數據傳輸的開銷就越大，副作用就越大。
為了提高GPU服務器的資源利用率，嘗試了一些加速的手段。
基於Pytorch1.6.0版本實現，官方支持
amp功能，不再需要外部apex庫；
此外比較重要的庫是Dali。
梳理了訓練框架，並將參考代碼放到Github上。
如果覺得對你有所啟發，請給個star呀。
1. 訓練速度的瓶頸及應對思路
這邊主要說的是CV領域，但在其他領域，思路應該也是相通的。
模型訓練過程中，影響整體速度的因素主要有以下幾點：
- 將數據從磁盤讀取到內存的效率；
- 對圖片進行解碼的效率；
- 對樣本進行在線增強的效率；
- 網絡前向/反向傳播和Loss計算的效率；
針對這幾個因素，分別采取如下幾種應對思路：
- 加快數據讀取可以有幾種思路：
- 采取類似TF的tfrecord或者Caffe的lmdb格式，提前將數據打包，比反復加載海量的小文件要快很多，但pytorch沒有通用的數據打包方式；
- 在初始化時，提前將所有數據加載到內存中（前提是數據集不能太大，內存能裝得下）；
- 將數據放在SSD而非HDD，可以大大提速（前提是你有足夠大的SSD）；
- 提升圖片解碼速度，可以考慮采用NVIDIA-DALI庫，能夠利用GPU來加速JPG格式的圖片解碼，針對其他格式的圖片（如PNG），不能實現GPU加速，但也可以兼容；
- 提升樣本在線增強的效率，同樣可以通過NVIDIA-DALI庫，實現GPU加速；
- 在網絡結構確定的情況下，提速主要有兩種方式，並且可以同時采用：
- 采用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占用
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資源，不會造成硬件資源的浪費；
- 通過
Nvidia-Dali模塊的合理配置，可以顯著提升數據加載和在線增強階段的效率，特別是在訓練一些輕量化模型時，往往瓶頸不在於GPU的計算速度，而在於CPU等其他部件的負載； - 通過
DistributedDataParallel模塊的合理配置，可以實現多卡的負載均衡，不論是顯存占用還是GPU利用率，都能夠達到平衡，不會有其中1張卡變成效率瓶頸； - 通過
torch.cuda.amp模塊的合理配置，可以進一步降低顯存占用，從而可以設置更大的batch_size，提高模型收斂速度； torch.cuda.amp模塊還可以顯著降低網絡前向推理時間，從而進一步提高訓練效率。
綜合應用如上所述的手段，基本上可以實現顯卡數量和訓練效率之間的線性增長關系。
不會發生卡多了，但是單卡的效率卻大大下降的現象。
6. 一些意外
原以為本篇到此就該結束了，但又遇到了新的問題。
當訓練執行一段時間后，由於整個系統長時間處於高負載的狀態，顯卡溫度飆升，觸發了顯卡的保護機制，自動降頻了，GPU利用率直接降到了原來的一半左右。
之前顯卡運行效率低的時候，散熱不良的問題還沒有顯露出來，一旦長時間高負荷運轉，多卡密集排布和風冷散熱的不足就暴露出來了。
下一步是要折騰水冷散熱了么？