torch單機多卡重點：

 

針對於單機多卡，可以使用 nn.DataParallel 進行訓練，需要注意的是，與單卡訓練有一些地方不同：

（1）模型和優化器的初始化：

模型初始化之后，需要傳入 nn.DataParallel ，再進行並行化處理，同時注意優化器同樣需要做並行化處理，如下：

model = Model.Model(param).cuda()

model = nn.DataParallel(model, device_ids=device_ids) # multi-GPU 

optimizer = torch.optim.Adam(model.parameters(), lr=0.002)  # 優化器

optimizer = nn.DataParallel(optimizer, device_ids=device_ids) # multi-GPU 

 

device_ids是一個list，保存要使用的GPU標號：

torch.cuda.set_device(1)

device_ids = [1, 2]

 

（2）優化器參數的更新：

optimizer.module.step()

與常規相比，多了個.module，是因為在並行化處理的時候，torch也對模型參數的結構進行了外部的封裝，因此需要多一層才能獲取到參數

 

（3）模型的保存：

torch建議只保存參數，可移植性更強：

torch.save(model.module.state_dict(), root_path+'data/model/train: epoch:'+str(epoch+1)+',step:'+str(step)+',loss:'+str(min_train_loss)[:6]+'.pth')

還有就是模型整個保存，官方並不推薦，先不展示了

 

（4）模型的加載：

這里只展示保存的參數的加載，也就是將上面方法保存的模型進行加載：

# 加載模型Model
model = Model.Model(param)
model = nn.DataParallel(model, device_ids=[1, 2]).cuda() # multi-GPU 
# 獲得模型參數
model_dict = torch.load('/home/xxx/project/xxx/data/model/train: epoch:1,step:266,loss:1.4726.pth')
# 載入參數
model.module.load_state_dict(model_dict)

 

其他：

 nn.DataParallel  可以實現單機多卡的訓練，但是模型的保存也分為很多種，比如單機多卡訓練出的模型，如何在無GPU的環境下加載，如何在單卡情況下加載等等…

然后，還有一些分布式訓練（多機多卡）的方法也可以用在單機多卡的情況下，這些還沒有進行嘗試。