PyTorch测试时候显存爆炸问题

 

问题描述

CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 2.00 GiB total capacity; 1.13 GiB already allocated; 0 bytes free; 1.15 GiB reserved in total by PyTorch)

 

猜测：测试时候未有释放显卡内存，导致每次加载模型，显卡内存都会爆炸，就很奇怪，明明测试时候只预测后面一个数据。

找了相关的文章，发现可以通过单步调试，查看显存变化的方法。原来pytorch会把每次的梯度保存在显存中，需要自行清空。

通过 with torch.no_grad():

语句可以不积累梯度图。

测试结果，使用如下代码

with torch.no_grad():

            output,hidden_state = model(x)

不再提示显卡内存不足的情况

总结

• 梯度图内存需要自行清理

参考：https://blog.csdn.net/dong_liuqi/article/details/119239671