pytorch多gpu多模型训练时报错AssertionError: Invalid device id

初入pytorch,机器上有两块gpu,分别同时训练两个模型,分别添加了如下代码加以区分:

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

第一块gpu上的模型正常运行,但在第二块gpu上的模型运行时报错 AssertionError: Invalid device id,定位错误代码为:

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

model = DFL_VGG16(k=10, nclass=args.nclass)
device = torch.device("cuda:1")
model = model.to(device)    # 定位到这一行

于是进行了如下测试:

import torch
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"    
print(torch.cuda.device_count())            # 1
print(torch.cuda.current_device())          # 0
print(torch.cuda.get_device_name())         # GeForce RTX 2080 Ti

os.environ["CUDA_VISIBLE_DEVICES"] = "1"    
print(torch.cuda.device_count())            # 1
print(torch.cuda.current_device())          # 0
print(torch.cuda.get_device_name())         # GeForce RTX 2080 Ti



即表明如果指定了某一块gpu,不管指定的是第几块,它的名字都为"cuda:0"

所以将错误代码进行如下修改,代码就正常运行了!

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

model = DFL_VGG16(k=10, nclass=args.nclass)
device = torch.device("cuda:0")    # 不管指定的是第几块gpu,gpu的id都为0
model = model.to(device)

 

你可能感兴趣的:(pytorch)