弄清Pytorch显存的分配机制
作者:颀周 时间:2023-05-22 22:12:44
对于显存不充足的炼丹研究者来说,弄清楚Pytorch显存的分配机制是很有必要的。下面直接通过实验来推出Pytorch显存的分配过程。
实验实验代码如下:
import torch
from torch import cuda
x = torch.zeros([3,1024,1024,256],requires_grad=True,device='cuda')
print("1", cuda.memory_allocated()/1024**2)
y = 5 * x
print("2", cuda.memory_allocated()/1024**2)
torch.mean(y).backward()
print("3", cuda.memory_allocated()/1024**2)
print(cuda.memory_summary())
输出如下:
代码首先分配3GB的显存创建变量x,然后计算y,再用y进行反向传播。可以看到,创建x后与计算y后分别占显存3GB与6GB,这是合理的。另外,后面通过backward(),计算出x.grad,占存与x一致,所以最终一共占有显存9GB,这也是合理的。但是,输出显示了显存的峰值为12GB,这多出的3GB是怎么来的呢?首先画出计算图:
下面通过列表的形式来模拟Pytorch在运算时分配显存的过程:
如上所示,由于需要保存反向传播以前所有前向传播的中间变量,所以有了12GB的峰值占存。
我们可以不存储计算图中的非叶子结点,达到节省显存的目的,即可以把上面的代码中的y=5*x与mean(y)写成一步:
import torch
from torch import cuda
x = torch.zeros([3,1024,1024,256],requires_grad=True,device='cuda')
print("1", cuda.memory_allocated()/1024**2)
torch.mean(5*x).backward()
print("2", cuda.memory_allocated()/1024**2)
print(cuda.memory_summary())
占显存量减少了3GB:
来源:https://www.cnblogs.com/qizhou/p/14110086.html
标签:Pytorch,显存,分配
![](/images/zang.png)
![](/images/jiucuo.png)
猜你喜欢
go语言制作的zip压缩程序
2024-02-12 21:03:00
在Python反编译中批量pyc转 py的实现代码
2023-08-11 20:21:04
![](https://img.aspxhome.com/file/2023/2/61732_0s.png)
XML简易教程之一
2008-09-05 17:19:00
PHP 危险函数全解析
2023-11-20 15:36:12
python常规方法实现数组的全排列
2021-11-20 23:09:29
python+splinter实现12306网站刷票并自动购票流程
2023-08-02 20:50:18
Python多线程中阻塞(join)与锁(Lock)使用误区解析
2022-03-22 08:00:31
如何将PySpark导入Python的放实现(2种)
2022-10-21 02:12:16
python装饰器相当于函数的调用方式
2021-05-13 13:39:03
asp防范SQL注入攻击的函数
2008-03-11 12:23:00
Django中数据在前后端传递的方式之表单、JSON与ajax
2022-08-09 11:32:44
Golang语言如何高效拼接字符串详解
2023-07-11 03:46:57
TypeScript学习之强制类型的转换
2024-04-25 13:08:07
MySQL多表查询与7种JOINS的实现举例
2024-01-12 23:34:32
![](https://img.aspxhome.com/file/2023/1/68581_0s.png)
仅用50行Python代码实现一个简单的代理服务器
2022-11-26 14:44:57
Windows上配置Emacs来开发Python及用Python扩展Emacs
2022-07-20 01:55:47
Python+OpenCV图像处理—— 色彩空间转换
2022-02-12 16:28:01
![](https://img.aspxhome.com/file/2023/0/126940_0s.png)
js序列化和反序列化的使用讲解
2023-08-05 08:13:24
django 使用内置messages的操作
2023-04-23 13:35:56
Python学习之直方图均衡化原理详解
2022-03-06 19:08:07
![](https://img.aspxhome.com/file/2023/2/82772_0s.png)