用GPU跑深度学习的时候显存溢出

2022.10.29 本文记录一次跑深度学习时显存溢出的情况

在本机(windows、NVIDIA GeForce RTX 3060、6G显存)跑图算法时显存溢出,如下:

在这里插入图片描述
这里显示我的显存总共6G,跑该代码需要3.45G,但留给PyTorch的只剩3.65G,比所需要的并未超出多少,所以会出现显存溢出的情况。

我的解决办法是:

  • 首先,缩小batch_size,原来是10,现在缩小为5来进行训练:

在这里插入图片描述

在这里插入图片描述

  • 然后,我在每个epoch训练后加入清理缓存的操作,即
torch.cuda.empty_cache()

在这里插入图片描述

做完上述两步即可正常运行了。

另外,要注意的一点是,在test或者validate的时候,是不需要计算存储梯度参数值的,所以记得使用:with torch.no_grad():循环

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐