pytorch 多gpu訓練

pytorch 多gpu訓練

用nn.DataParallel重新包裝一下

數據並行有三種情況

前向過程

device_ids=[0, 1, 2]
model = model.cuda(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)

只要將model重新包裝一下就可以。

后向過程

optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.001)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)
#因為它在DataParallel里面，所以要先變成普通的nn.SGD對象，然后才能調用該類的梯度更新方法。
optimizer.module.step() 

在網上看到別人這樣寫了，做了一下測試。但是顯存沒有變化，不知道它的影響是怎樣的。
更新學習率的時候也需要注意一下：

for param_lr in optimizer.module.param_groups: #同樣是要加module
    param_lr['lr'] /= 2

criterion(loss 函數)

def init_criterion():
    criterion = loss.CrossEntropyLoss2d()
    criterion = torch.nn.DataParallel(
            criterion, range(gpu_nums)).cuda()  # range(self.settings.n_gpu)
    return criterion
    
# criterion = init_criterion()
criterion = loss.CrossEntropyLoss2d()

這個並行的效果對顯存是有影響的，但是效果不明顯。我沒有做太多實驗。
訓練的時候會出現問題：

loss = criterion(out, labels_tensor)
loss /= N
optimizer.zero_grad()
# loss.backward()
loss.sum().backward()

數據並行返回的結果的維度和之前維度是不一樣的所以反向傳播的時候需要做一下修改