【CUDA OUT OF MEMORY】【Pytorch】计算图与CUDA OOM

计算图与CUDA OOM

在实践过程中多次碰到了CUDA OOM的问题,有时候这个问题是很好解决的,有时候DEBUG一整天还是头皮发麻。

最近实践对由于计算图积累导致CUDA OOM有一点新的看法,写下来记录一下。包括对计算图的一些看法和一个由于计算图引发错误的简化实例记录。

本人能力有限,认识片面如果犯了错误希望大家指教!

计算图的存储

计算图是pytorch进行梯度反向传播核心,计算图是在程序运行过程中动态产生的,当tensor变量赋予了requires_grad=True的属性时,torch会自动记录其参与的计算并形成计算图保存在显存中。

敲重点:计算图是会吃显存的! 本来想截下来描述一下计算图是长什么样的,至少是概念的表述一下,结果去学习了一圈发现:和我想的完全不一样!附上学习链接:传送门。更关键的是我还没完全看懂学会(🐶),有没有大大学会了教我一下,不甚感激!

总的来说一个tensor它内部包含的grad_fn别有洞天,首先grad_fn也是作为一个节点在计算图中的(其在pytorch的C艹中是Node的子类),grad_fn不仅是记录了这个tensor是被什么数学符号计算来的,它还暗搓搓记录了这个tensor是是从哪些数字里头窜出来的,以及其和其他grad_fn的py友谊,还有被包含在其内部context中的信息,我偷那个学习链接的一张图展示一下一个计算图的形态,借花献佛,展示一下grad_fn偷偷摸摸用你的卡干了啥事情。

BTW,提几个小知识点

  • 我们常用的detach()方法,就是通过把tensor的grad_fn扬了从而把tensor从计算图中剥离出来。

    x
    tensor([1.], requires_grad=True)
    y = x+1
    y.grad_fn
    <AddBackward0 object at 0x7f8306e68b50>
    y.detach().grad_fn is None
    True

  • 关于*.backward(retain_graph=True)的问题,backwardretain_graph默认是False,其含义是经过默认的*.backward()之后,计算图会被清空从而释放其占用的显存。和detach不一样的是,grad_fn还是那个grad_fn只不过它悄咪咪维持的友谊被杀掉了,如下:

    x
    tensor([1.], requires_grad=True)
    y = x+1; y.grad_fn
    <AddBackward0 object at 0x7f8306e68b50>
    y.backward(retain_graph=False)
    y.grad_fn
    <AddBackward0 object at 0x7f8306e68b50>

  • 续上面一点的内容,但是内容包含我瞎猜的成分(🐶),我们猜测一下backward杀掉了grad_fn的什么东西。一般的,我们认为当retain_graph=False的时候,我们只能backward()一次,因为计算图会被清空,第二次尝试反向传播会造成错误。但其实不然!如下实验例子1的尝试,我们连续backwrad并没有报错。AMAZING啊!。进一步的我们进行例子2的实验,我们只是简单的让前向多了一个乘法计算,然后另z反向传播两次,这回顺理成章的报错,同时报错之后我们再次反传y,我们发现反传y又不会报错。我猜测:backward()会清楚grad_fn节点和其他grad_fn的联系,因此zgrad_fn不能联系到ygrad_fn了,于是第二次z.backward()报错,但是y直接和叶子x连接,不需要其他的grad_fn朋友也能自己和自己玩。

    例子1:

    x
    tensor([1.], requires_grad=True)
    y = x+1
    y.backward(retain_graph=False);y.backward()
    返回没有报错!


    例子2:

    x
    tensor([1.], requires_grad=True)
    y=x+1;z=2*y #前向过程多了一个乘法
    z.backward(retain_graph=False)
    z.backward()
    Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
    File "/Users//opt/anaconda3/lib/python3.8/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
    File "/Users/
    /opt/anaconda3/lib/python3.8/site-packages/torch/autograd/init.py", line 200, in backward
    Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
    RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
    y.backward()
    返回没有报错

一个由于没处理好计算图导致OOM的例子

python 复制代码
import torch,time
l1 = torch.nn.Linear(1000,1000).cuda()
l2 = torch.nn.Linear(1000,1000).cuda()
memory = []

for _ in range(10000000):
	time.sleep(0.01)
	data_input = torch.rand(1000).cuda()
	output = l1(l2(data_input))
	output.backward(retain_graph=True) #此行与报错无关 
	memroy.append(output.cpu()) 
	#memory存储的内容通过.cpu()转移在主存上,
	#但是与output相关联的l1,l2的计算图依旧停留在显存中,并在循环中一直积累撑爆显存。
	...some other operations...

这个例子中,由于每个output不能被正常清除计算图显存,最终导致OOM。

这个例子是某次实践的超级简化版,如果只看这个例子的话,其实只要把最后一行改写成

python 复制代码
memory.append(output.detach().cpu())

就会由于output在每次循环后失去引用(detach()创建了新的变量)从而被回收,计算图被自动清空避免OOM。

相关推荐
微学AI1 分钟前
融合注意力机制和BiGRU的电力领域发电量预测项目研究,并给出相关代码
人工智能·深度学习·自然语言处理·注意力机制·bigru
知来者逆13 分钟前
计算机视觉——速度与精度的完美结合的实时目标检测算法RF-DETR详解
图像处理·人工智能·深度学习·算法·目标检测·计算机视觉·rf-detr
一勺汤16 分钟前
YOLOv11改进-双Backbone架构:利用双backbone提高yolo11目标检测的精度
人工智能·yolo·双backbone·double backbone·yolo11 backbone·yolo 双backbone
武汉唯众智创18 分钟前
高职人工智能技术应用专业(计算机视觉方向)实训室解决方案
人工智能·计算机视觉·人工智能实训室·计算机视觉实训室·人工智能计算机视觉实训室
Johny_Zhao29 分钟前
MySQL 高可用集群搭建部署
linux·人工智能·mysql·信息安全·云计算·shell·yum源·系统运维·itsm
拾忆-eleven33 分钟前
C语言实战:用Pygame打造高难度水果消消乐游戏
c语言·python·pygame
一只可爱的小猴子39 分钟前
2022李宏毅老师机器学习课程笔记
人工智能·笔记·机器学习
地瓜机器人44 分钟前
乐聚机器人与地瓜机器人达成战略合作,联合发布Aelos Embodied具身智能
人工智能·机器人
带娃的IT创业者1 小时前
《AI大模型趣味实战》基于RAG向量数据库的知识库AI问答助手设计与实现
数据库·人工智能
旦莫1 小时前
Python 教程:我们可以给 Python 文件起中文名吗?
开发语言·python