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上，我一般报错都是因为漏掉了其中一两个。