Pytorch使用torch.nn.DataParallel()在Linux服务器上多GPU训练

前期我自己也在网上也找了很多相关的文章，但是都不能用。后来结合Pytorch中多GPU训练这篇文章，加上自己的尝试，终于成功在服务器上使用多gpu训练成功。

首先，使用代码

os.environ["CUDA_VISIBLE_DEVICES"] = "2， 3， 4" 

这句代码能够告诉程序接下来要使用的gpu编号是哪几块。比如服务器上有8块gpu，编号为0，1， 2，……，7，你只想用其中的2,3,4块的话，就可以使用上述代码。

接下来你需要使用代码

device = torch.device('cuda')

告诉程序你使用的是gpu而不是cpu训练模型.只用'cuda'默认的是'cuda:0'，也就是前边os.environ语句指定的第一块gpu，编号为2。你也可以改为device = torch.device('cuda:1')，这样指定的gpu编号为3。
这里指定的“device”是你多gpu训练的主gpu，它所占用的显存会比其他gpu的显存高。

然后你就可以使用代码

model = torch.nn.DataParallel(Model(), device_ids=[0,1，2])
# 将模型送入gpu
model.to(device)
# 将损失函数送入gpu
loss = nn.MSELoss().to(device)
# 将数据和标签送入gpu
images = images.to(device)
label= label.to(device)

就可以在多gpu上训练你的模型了。

注意： torch.nn.DataParallel(Model(), device_ids=[0, 1, 2])里边的device_ids，0对应前边的gpu编号2，1对应gpu编号3，2对应gpu编号4，以此类推。