[原创]关于解决pytorch训练神经网络时显存一直增长的问题

@[原创]关于解决pytorch训练神经网络时显存一直增长的问题

问题描述

在训练自定义loss和自定义网络结构的一个模型的时候，发现模型和数据都比较简单的情况下，在训练过程中居然把24g的显卡拉爆了显存。

然后使用nvidia-smi -l观察显存变化，发现是有规律的显存一直增加，直到OOM。

问题解决思路

在这个过程中尝试询问了chatgpt，但是发现它提供的解决方案，诸如torch.cuda.memory_cached() /del data 等命令放在每次循环后面并不能解决问题。

所以后面尝试在谷歌进行搜索，找到了下面这篇的知乎的博客：

链接: link.

这篇文章的四种方法其实都没有解决我的问题，但是它的第一种情况给了我一点启发，此外chatgpt在最开始提到的原因也是关键，促成了后面问题的解决。

在直接尝试各种解决手段无果后，我决定自行去查看代码段和显存的使用情况，主要是使用下面这个命令：

print("Memory Allocated:", torch.cuda.memory_allocated() / (1024 ** 2), "MB")

这个相当于代码内部的nvidia-smi可以查看代码运行到这里的时候显存的占用量（这里说的并不严谨，因为还有memory_cached这个命令查看缓存的显存，但是先这么肤浅的使用吧。）

随后我发现了显存的一直增长来自于下面这个语句：

self._update_stats({'Loss/total': loss_iter.item(), 'PSNR/initial': psnr_initial}, batch_size, loader)

回想起之前的那篇知乎博客的第一点解决方案，我把psnr_initial改为了psnr_initial.item()，随后发现占用的内存不会再随着循环的迭代而增长了。

分析

总结来说在_update_stats()这个函数中，它会使用psnr_initial这个tensor，所以如果不适用item()做隔离的话，相当于gpu上有数据一直在被其他函数调用，所以主循环里每一次迭代都无法释放上一次的显存，但是用了item()就可以跟gpu上的数据隔离开，帮助主循环迭代过程的显存正常释放。

所以在遇到同样的问题的时候，朋友们可以主要检查一下每次迭代结束是否有gpu上的数据仍然在被其他的函数调用，比如loss的记录函数。