RuntimeError: module must have its parameters and buffers on device cuda:2 (device_ids[0]) but found one of them on device: cuda:0
出现这个问题是使用一机多卡分布式训练时,没有调用显卡0,例如使用显卡1,2,3进行训练。
使用的模块是pytorch nn.DataParallel。
下面给出有问题的代码
device = torch.device("cuda" if args.cuda else "cpu")
#Try to load models
if args.model == 'pointnet':
model = PointNet(args).to(device) #这里应该对应自己的模型
elif args.model == 'dgcnn':
model = DGCNN(args).to(device) #这里应该对应自己的模型
else:
raise Exception("Not implemented")
device_ids = [1,2,3]
model = nn.DataParallel(model,device_ids=device_ids)
问题就出在device处,因为下面这条语句,默认将模型放在了cuda 0上,也就是我们没有使用的显卡0。于是会产生错误RuntimeError: module must have its parameters and buffers on device cuda:2 (device_ids[0]) but found one of them on device: cuda:0
device = torch.device("cuda")
那么修改的思路也比较简单,如下。指定cuda,避免将其放在默认的显卡0上。
device = torch.device("cuda:1" if args.cuda else "cpu")
#Try to load models
if args.model == 'pointnet':
model = PointNet(args).to(device) #这里应该对应自己的模型
elif args.model == 'dgcnn':
model = DGCNN(args).to(device) #这里应该对应自己的模型
else:
raise Exception("Not implemented")
device_ids = [1,2,3]
model = nn.DataParallel(model,device_ids=device_ids)
问题解决,可以愉快的训练了。