pytorch multi-gpu train

記錄一下pytorch如何進行單機多卡訓練：

官網例程：https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html

下面以一個例子講解一下，例如現在總共有8張卡，在第5、6、7三張卡上進行訓練；

step 1:可視化需要用到的GPU

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "5 , 6 , 7"

device = torch.device("cuda:0")  #注意多卡訓練的時候，默認都先將model和data先保存在id:0的卡上（即實際的第5塊卡），然后model的參數會復制共享到其他卡上，data也會平分成若干個batch到其他卡上（所以第一塊卡上稍微耗費一點顯存）；

device_ids = [0 , 1 , 2] #注意device_ids必須從0開始，代碼中的所有的device id都需要從0開始（這里的0代表第5塊卡，1代表第6塊卡，類推）；

step 2:利用DataParallel對Model類進行封裝

model = nn.DataParallel(model , device_ids = device_ids)

model.to(device)

step 3:

data.to(device)  #id:0卡上的數據再被平分成若干個batch到其他卡上

注意：晚上還有一些例程，需要對optimizer和loss利用DataParellel進行封裝，沒有試驗過，但上面方法是參考官網例程，並經過實操考驗；