Pytorch RuntimeError: all tensors must be on devices[0]

问题:
在跑多GPU时出现RuntimeError: all tensors must be on devices[0]问题,找了一番没解决,最后瞎试出来了。
主体代码部分:

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
net.cuda()
net = torch.nn.DataParallel(net)

第一次训练时设置的args.gpu=”0,2”,第二次同样配置就出现RuntimeError: all tensors must be on devices[0]问题。google一番,找到几个解决办法:
办法1:
CUDA_VISIBLE_DEVICES中gpu id和DataParallel中device_ids不同,如果使用:

CUDA_VISIBLE_DEVICES = 3,1,2

则:

net = torch.nn.DataParallel(net, device_ids=[0,1,2])

意思是此处的device_ids=[0,1,2]分别对应的是三个gpu3,1,2。但我使用无效。
参考链接
办法2:
使用net.cuda(0)或者with torch.cuda.device(0):指定数据载入gpu。又对我没用。
参考链接
办法3:
在device_ids中只指定一个gpu,代码跑的时候还是使用两个gpu。

os.environ['CUDA_VISIBLE_DEVICES'] = "0,2"
net.cuda()
net = torch.nn.DataParallel(net, device_ids=[0])

猜你喜欢

转载自blog.csdn.net/qq_16234613/article/details/80948032