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()
数据并行返回的结果的维度和之前维度是不一样的所以反向传播的时候需要做一下修改
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:pytorch 多gpu训练 - Python技术站