PyTorch指定GPU並行訓練小記

假設一臺機器上有多塊GPU。給定需要訓練的模型,每塊GPU將分別獨立維護一份完整的模型參數。在模型訓練的任意一次迭代中,給定一個小批量,將該批量中的樣本劃分成若干份並分給每塊GPU一份。然後,每塊GPU將分別根據自己分到的訓練數據樣本和自己維護的模型參數計算模型參數的梯度。接下來把這若干塊GPU上分別計算得到的梯度相加,從而得到當前的小批量梯度。之後,每塊GPU都使用這個小批量梯度分別更新自己維護的那一份完整的模型參數。

在多GPU的機器上,要選擇特定的GPU來運行程序,可在程序運行命令前使用命令:

CUDA_VISIBLE_DEVICES=0

在這裏插入圖片描述
其中0爲服務器中的GPU編號,例如在有8塊GPU的機器上可以爲0,1,2,3,4,5,6,7等,編號可以通過以下命令查看:

nvidia-smi

通過指定GPU編號表明對當前啓動多程序或者腳本可見的GPU編號,未指定的GPU則不會被該程序或腳本使用。

需要注意,由CUDA_VISIBLE_DEVICES指定的GPU編號和代碼中索引的GPU編號不完全對等,具體可見下面的例子:

CUDA_VISIBLE_DEVICES=1     # 只有編號爲1的GPU對程序是可見的,在代碼中gpu[0]指的就是這塊兒GPU
CUDA_VISIBLE_DEVICES=0,2,3 # 只有編號爲0,2,3的GPU對程序是可見的,在代碼中gpu[0]指的是第0塊兒,
						   # gpu[1]指的是第2塊兒,gpu[2]指的是第3塊兒
CUDA_VISIBLE_DEVICES=2,0,3 # 只有編號爲0,2,3的GPU對程序是可見的,但是在代碼中gpu[0]指的是第2塊兒,
						   # gpu[1]指的是第0塊兒,gpu[2]指的是第3塊兒
torch.device(f"cuda:{args.local_rank}")#args.local_rank由PyTorch自己根據GPU使用情況指定

torch.device("cuda:0")#gpu[0]

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
CUDA_VISIBLE_DEVICES=0,1,2,3 python gpu_test.py

參考資料:
[1] 『MXNet』第七彈_多GPU並行程序設計
[2] CUDA GPU 隔離 CUDA_VISIBLE_DEVICES

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章