pytorch解决两个GPU同时训练问题。

解决两个GPU同时训练问题。

使用场景

---

我有两个GPU卡。我希望我两个GPU能并行运行两个网络模型。

代码

错误代码1：

#对于0号GPU
os.environ['CUDA_VISIBLE_DEVICES']='0,1'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#对于1号GPU
os.environ['CUDA_VISIBLE_DEVICES']='0,1'
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

0号GPU不报错，1号GPU报错。错误如下
RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 0 does not equal 1 (while checking arguments for cudnn_convolution)

错误代码2：

#对于0号GPU
os.environ['CUDA_VISIBLE_DEVICES']='0'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#对于1号GPU
os.environ['CUDA_VISIBLE_DEVICES']='1'
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

0号GPU不报错，1号GPU报错。错误如下
CUDA: invalid device ordinal
正确代码如下：

#对于0号GPU
os.environ['CUDA_VISIBLE_DEVICES']='0'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#对于1号GPU
os.environ['CUDA_VISIBLE_DEVICES']='1'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")