Pytorch 分佈式訓練技巧
這是一篇統整了 Pytorch DistributedDataParallel, Apex, warm-up, learning rate scheduler 用法的文章,也會提到 early-stopping 和 Random seed 的設置。
5 min readAug 25, 2020
在下面的章節中我將依序介紹怎麼使用 DistributedDataParallel (DDP)、如何將 Apex, warmup, learning rate scheduler 這三個訓練技巧融入 DDP 的訓練中,也會提到 early-stopping 和 Random seed 的設置,完整的程式碼置於 GitHub。
If you need to read this article in English, please check my GitHub link below:
https://github.com/Lance0218/Pytorch-DistributedDataParallel-Training-Tricks
DistributedDataParallel (DDP)
- Pytorch 官網也建議在使用多 GPU 訓練時可以用 DistributedDataParallel (多進程控制多GPU) 取代使用 DataParallel (單進程控制多 GPU),在提升速度的同時也解決了 GPU 負載不平均的問題。
- 基本用法就是將要使用的 model 載入,並用 DDP 包起來 (L16),local_rank 是調用
torch.distributed.launch
產生的當前 GPU 的 rank。 - 若是在多節點下傳輸資料會嚴重影響效率,故 DistributedSampler 用來確保 DataLoader 只會載入數據集的特定子集 (L4),在 DistributedSampler 使用下的 batch_size 為單張 GPU 實際使用的 batch 尺寸。
- 需要在每個 epoch 創造 DataLoader 迭代前就呼叫 set_epoch(epoch) 來確保 shuffle 的正常運作 (L21)。
Apex
- Apex 透過犧牲一些精度來減少顯存佔用,讓我們可以使用更大的網路架構和更大的 batch size。
- 透過
amp.initialize
來初始化你建構好的模型和優化器 (L12),如果像我一樣有使用非 Pytorch 內建的優化器 (比如 lookahead),則要注意初始化的對象。 - 並用以清空過往梯度 (L19)、反向傳播計算當前梯度 (L23)、根據當前梯度更新網路參數 (L24)。
- 若要使用 sigmoid, softmax 這類函數 (詳情請見
apex.amp
的 GitHub),則需要在初始化前先註冊 (L11)。
Warmup & learning rate scheduler & early-stopping
- Warmup 是一種訓練技巧,透過由小到大預熱學習率可以避免一開始學習率過大所造成的不穩定,並可以更快找到穩定優化的方向,已經在一些任務上被證明有效,本篇透過 lr_scheduler 實現 (L6, L7, L28, L29)。
- Learning rate scheduler 也是訓練模型的技巧,本篇以我比較喜歡用的
lr_scheduler.ReduceLROnPlateau
來作例子 (L8, L30),要注意lr_scheduler
裡面的優化器要指向 Pytorch 內建的而不是自己額外加的 (L7, L8)。 - 本篇也有加入 early-stopping 來減少過擬合 (L9, L31~L33),另外還有用 consine learning rate 做衰減的例子,有興趣的請在此處查看。
- 本篇的寫法 warmup 學習率是從 0 開始,若想一開始就有值,可以加上 (L12, L13)、用 (L17, L18) 取代 (L28, L29)。
Random seed
- 若要訓練可以重複實現的模型則要設定好隨機種子。
- 需要設定的地方有兩個:Training DataLoader 前 (L5) 用來固定 shuffle 結果、model 前 (L11) 用來固定起始權重。
- 本篇我直接用函式 same_seeds 一次調整所有的隨機種子,你也可以嘗試只調整特定的隨機種子。
Run
- Pytorch 官方提供兩種運行方法:
torch.distributed.launch
和torch.multiprocessing.spawn
,但我使用後者有時會有訓練結束後 GPU 不會自動釋放的狀況,所以本篇使用torch.distributed.launch
來做演示。 - 本篇主要演示單節點多 GPU 運行方式:CUDA_VISIBLE_DEVICES 指定使用的 GPU, nproc_per_node 為該節點使用的 GPU 數,若需要執行多個訓練將 master_port 錯開即可,這裡有更詳細的說明。
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --master_port=6666 DDP_warmup.py -en=DDP_warmup
本篇運行環境為 CUDA10.2, cuDNN7.6.5, Python3.7, Pytorch1.3.1
Function implementation
- Lookahead 實現:lookahead.pytorch
- Early-stopping, random seed 實現:customized_function.py
感謝觀看,有任何問題歡迎留言討論或連繫我,祝有個愉快的一天😉