PyTorch训练显存越跑越涨，隐式保留计算图导致OOM，怎么办？🔍

你有没有遇到过这种情况：PyTorch训练模型时 一开始显存占用还行，但越跑越卡，再说说直接OOM？ 是不是？ 你是不是也像我一样，盯着屏幕上的“CUDA out of memory”报错，心里一阵阵发毛？

显存越跑越高，这到底是谁的锅？

蚌埠住了... 其实这事儿不完全是PyTorch的错，但它的“锅”也不小。问题出在PyTorch的“计算图”上。你可能没意识到，PyTorch默认会为每个张量保留完整的计算图，哪怕你只是想看看模型输出个啥。这就像你去超市买东西，后来啊收银员非得把每样商品的生产记录都给你背一遍，你说烦不烦？

不错。 而更烦的是这些记录还赖在显存里不走，哪怕你已经用不上它了。于是乎，显存就像吹气球一样，越吹越大，再说说“砰”地一声炸了。

计算图是个啥玩意儿？

简单说计算图就是记录你模型是怎么一步步算出来的。比如你输入x，经过一层层网络，再说说输出y。PyTorch会记住每一步的计算过程，方便你回头求导。但问题是它记住了就不放，哪怕你已经用不着了它还死死攥着不放。这就像你吃完饭，服务员还非得看着你把盘子舔干净，你说气不气，ICU你。？

所以 显存越跑越高，不是你代码写错了而是PyTorch太“记仇”了啥都给你记着，哪怕你已经不想再见到它了，基本上...。

显存爆炸的“罪魁祸首”

我们来列几个“罪魁”：

• retain_graph=True这个参数就像一个“粘人精”， 你一用它，PyTorch就死活不释放计算图，哪怕你已经跑完一个batch了它还赖着不走。
• loss += loss你以为这只是个累加？错！它会把每次的loss都保留下来然后越加越多，再说说显存直接被撑爆。