pytorch：多GPU训练、保存加载参数、测试

一、指定使用单块gpu训练

os.environ["CUDA_VISIBLE_DEVICES"] = '1'

二、使用pytorch进行多GPU训练的需要添加的代码：

#设置使用哪些显卡，我这里用所有可用的
device_ids=range(torch.cuda.device_count())
model = Transformer().cuda()
model = nn.DataParallel(model, device_ids=device_ids)

三、保存和加载参数

#加载参数
model.load_state_dict(torch.load('模型参数/两条130数值的测试用的模型参数.pkl'))
#保存参数
torch.save(model.state_dict(), '模型参数/两条130数值的测试用的模型参数.pkl')

四、多GPU加载参数后测试：

# Test
enc_inputs = next(iter(loader))
greedy_dec_input = greedy_decoder(model, enc_inputs[0].view(1, -1).cuda(), start_value=enc_inputs[0][-1])
predict = model.module(enc_inputs[0].view(1, -1).cuda(), greedy_dec_input)
print(predict)

会有如下报错：
AttributeError: ‘DataParallel’ object has no attribute ‘encoder’

因为加载参数时，model = nn.DataParallel(model, device_ids=device_ids)，model等于了nn.dataparallel,在操作原始模型的时候需要将model换为model.module