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

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

问题描述

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

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

问题解决思路

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

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

链接: link.

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

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

python 复制代码
print("Memory Allocated:", torch.cuda.memory_allocated() / (1024 ** 2), "MB")

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

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

python 复制代码
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的记录函数。

相关推荐
李元豪24 分钟前
【知足常乐ai笔记】机器人强化学习
人工智能·笔记·机器人
沫儿笙25 分钟前
焊接机器人智能节气装置
人工智能·机器人
MidJourney中文版25 分钟前
老年人与机器人玩具的情感连接
人工智能·机器人·语音识别
Codebee27 分钟前
AI驱动的低代码革命:解构与重塑开发范式
人工智能·低代码·代码规范
数据库安全28 分钟前
首批|美创智能数据安全分类分级平台获CCIA“网络安全新产品”
大数据·人工智能·web安全
Dymc29 分钟前
【目标检测之Ultralytics预测框颜色修改】
人工智能·yolo·目标检测·计算机视觉
HuashuiMu花水木1 小时前
PyTorch笔记3----------统计学相关函数
人工智能·pytorch·笔记
算家计算1 小时前
6 亿参数玩转 20 + 语言!OuteTTS-1.0-0.6B本地部署教程,轻量模型也能 hold 住跨语言合成
人工智能·开源
柠檬味拥抱1 小时前
USB‑C 式的工具联接:MCP 的模块化及通用标准探讨
人工智能
柠檬味拥抱1 小时前
模块化控制协议(MCP)在网络中增强智能体执行效率的研究
人工智能