2021-04-19：怎么分布式训练模型

怎么分布式训练模型

当一台电脑有多个GPU时，可以进行分布式训练

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
if torch.cuda.device_count() > 1:
	print('we have ',torch.cuda.device_count(),'GPUs')
	model = nn.DataParallel(model,device_ids=[0,1])#DataParallel会拆分

#保存模型：必须加module
torch.save(net.module.state_dict(), path)
#加载模型：
net.load_state_dict(torch.load(path))