torch 单机多卡运行

将torch模型使用单机多卡运行

步骤

第一步:设置运行环境

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"
ids = [0, 1, 2]

if torch.cuda.is_available():
    device = torch.device("cuda:0")
else:
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    device = torch.device("cpu")

第二步:模型初始化,并将模型放在多卡上

model = mymodel()
if torch.cuda.device_count() > 1:
	print("Use", torch.cuda.device_count(), 'gpus')
	model = torch.nn.DataParallel(model)
	model.to(device)

第三步:将loss的计算函数放在多卡上

crossEntropy = torch.nn.CrossEntropyLoss().cuda()
loss = crossEntropy(batch_pred, batch_target)

第四步:将输入数据放在多卡上

X_train, y_train = data
X_train, y_train = X_train.cuda(), y_train.cuda()

常见问题

没有报错,但也没有使用多卡

有可能是第三步漏掉了,尝试检查计算loss的函数是否放在cuda上了。

CUDA error: an illegal memory access was encountered

这种报错的本质就是模型model、输入数据(input_image、input_label)没有全部移动到GPU(cuda)上。
**温馨提示:**debug时一定要仔细检查是否每一个输入变量以及网络模型都移动到了GPU上,我一般报错都是因为漏掉了其中一两个。

上一篇:DCNv2编译过程报错:error: identifier "THCudaBlas_SgemmBatched" is undefined


下一篇:【PYTORCH】Pytorch与GPU相关大合集(含测试代码)