pytorch 多GPU训练

1
2
3
4
5
net = Net() #Net是自定义的一个网络结构类
device_ids = [2, 4, 5]
cudnn.benchmark = True
net = net.cuda(device_ids[0])
net = nn.DataParallel(net, device_ids=device_ids) #使用dataParallel重新包装一下

1
2
3
4
5
6
7
lr = 1e-2
momentum = 0.9
weight_decay = 1e-3
param = get_param(net, lr)
optimizer = optim.SGD(param, momentum=momentum, weight_decay=weight_decay) #准备pytorch中的随机梯度下降方法
loss = nn.MSEloss()
optimizer = nn.DataParallel(optimizer, device_ids=device_ids) #将optimizer放入dataparallel中。

1
2
3
4
5
6
7
8
9
img = Variable(img, requires_grad=True).cuda(device_ids[0])
# 输入图片数据
gt = Variable(gt_heatmap).cuda(device_ids[0]) #ground truth
predicted = net(img) # net是DataParallel对象,img 作为输入会将分为3份（bachsize/3），等3个并行计算结束后再以该轴组合再一起，predicted和img的shape是一样的。
l = loss(gt, predicted)
# compute gradient and do SGD step
optimizer.zero_grad()
l.backward() #在这儿使用optimizer的相应的对象。
optimizer.module.step() #因为它在DataParallel里面，所以要先变成普通的nn.SGD对象，然后才能调用该类的梯度更新方法。类似的，还有其他的一些需要注意的地方，看下面：

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

1
2
3
4
5
6
7
8
model_dict = net.state_dict()
vgg_19_key = list(vgg_19.keys())
model_key = list(model_dict.keys())
from collections import OrderedDict
vgg_dict = OrderedDict()
for i in range(param_num):
vgg_dict[model_key[i]] = vgg_19[vgg_19_key[i]]
model_dict.update(vgg_dict)

1
2
3
for item, value in saved_state.items():
name = '.'.join(item.split('.')[1:])
trans_param[name] = value