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

2022-10-22 14:45:43

解决两个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")
  • 作者：werdery
  • 原文链接：https://blog.csdn.net/werdery/article/details/106154294
    更新时间：2022-10-22 14:45:43