前期我自己也在网上也找了很多相关的文章,但是都不能用。后来结合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,以此类推。