用nn.DataParallel重新包装一下

数据并行有三种情况

前向过程

代码语言：javascript复制

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

只要将model重新包装一下就可以。

后向过程

代码语言：javascript复制

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() 

更新学习率的时候也需要注意一下：

代码语言：javascript复制

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

criterion(loss 函数)

代码语言：javascript复制

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()

这个并行的效果对显存是有影响的，但是效果不明显。 训练的时候会出现问题：

代码语言：javascript复制

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

数据并行返回的结果的维度和之前维度是不一样的所以反向传播的时候需要做一下修改

model 函数 数据

0 人点赞