[]关于解决pytorch训练神经网络时显存一直增长的问题-Toy模板网

这篇具有很好参考价值的文章主要介绍了[]关于解决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的记录函数。文章来源地址https://www.toymoban.com/news/detail-629871.html

到了这里，关于[]关于解决pytorch训练神经网络时显存一直增长的问题的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！