torch 单机多卡运行

2021-12-26 16:02:49 阅读：205 来源： 互联网

标签：多卡上单机 torch 多卡 train cuda device model

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

步骤
常见问题
- 没有报错，但也没有使用多卡
- CUDA error: an illegal memory access was encountered

步骤

第一步：设置运行环境

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

标签：多卡上,单机,torch,多卡,train,cuda,device,model
来源： https://blog.csdn.net/hfut_lf/article/details/122156232

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9