CUDA out of memory

文章讲述了在使用PyTorch进行模型测试时遇到CUDA内存不足的问题,原因是未正确管理hook,导致GPU内存无法释放。hook函数在提取层的输入输出和权重后,如果没有及时移除,会持续占用内存,最终引发显存爆炸。解决方案是使用完毕后立即移除注册的hook。
摘要由CSDN通过智能技术生成

记录一下
模型测试时出现了CUDA out of memory,也是很奇怪
原来是有数据在GPU中一直没有释放,每次调用测试代码时都会累计,久而久之导致了CUDA out of memory

pytorch的hook机制可能导致,显存爆炸,hook函数取出某一层的输入输出跟权重后,不可进行存储,修改等操作,这会造成hook不能回收,进而导致取出的输入输出权重都可能不被pytorch回收,所以模型的负担越来也大,最终导致显存爆炸。

代码中有

outputs = []
def hook(module, input, output):
    outputs.append(output)
model.layer1[-1].register_forward_hook(hook)
model.layer2[-1].register_forward_hook(hook)
model.layer3[-1].register_forward_hook(hook)

原来是hook机制捣的鬼
使用完后,给他们remove一下就OK了

outputs = []
def hook(module, input, output):
    outputs.append(output)
h1 = model.layer1[-1].register_forward_hook(hook)
h2 = model.layer2[-1].register_forward_hook(hook)
h3 = model.layer3[-1].register_forward_hook(hook)
...
h1.remove()
h2.remove()
h3.remove()

问题解决

  • 0
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值