首页 > 其他 > 详细

torch.load CUDA ERROR: out of memory

时间:2020-05-12 12:53:42      阅读:160      评论:0      收藏:0      [点我收藏+]

在pytorch中load模型,出现CUDA ERROR:out of memory,基本有下面三个原因:

CUDA显存不足

解决办法:换成显存更大的显卡

显卡被占用

当存在多个显卡时,pytorch默认使用0号显卡,当0号显卡被占用时,即使1号显卡空闲仍然会报错
解决办法:让其他显卡不可见,此时0号显卡为空余的1号显卡

import os

os.environ[‘CUDA_VISIBLE_DEVICE‘]=‘1‘

模型保存GPU不同

使用cuda:0训练保存的模型,如果直接使用cuda:1去load,也会报错
解决办法:model.load_state_dict(torch.load(‘trained.pth‘, map_location={‘cuda:0‘: ‘cuda:1‘}))

参考

torch.load CUDA ERROR: out of memory

原文:https://www.cnblogs.com/huanxifan/p/12875255.html

(0)
(0)
   
举报
评论 一句话评论(0
关于我们 - 联系我们 - 留言反馈 - 联系我们:wmxa8@hotmail.com
© 2014 bubuko.com 版权所有
打开技术之扣,分享程序人生!