弄清Pytorch显存的分配机制


Posted in Python onDecember 10, 2020

对于显存不充足的炼丹研究者来说,弄清楚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())

输出如下:

弄清Pytorch显存的分配机制

代码首先分配3GB的显存创建变量x,然后计算y,再用y进行反向传播。可以看到,创建x后与计算y后分别占显存3GB与6GB,这是合理的。另外,后面通过backward(),计算出x.grad,占存与x一致,所以最终一共占有显存9GB,这也是合理的。但是,输出显示了显存的峰值为12GB,这多出的3GB是怎么来的呢?首先画出计算图:

弄清Pytorch显存的分配机制

下面通过列表的形式来模拟Pytorch在运算时分配显存的过程:

弄清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:

弄清Pytorch显存的分配机制

以上就是弄清Pytorch显存的分配机制的详细内容,更多关于Pytorch 显存分配的资料请关注三水点靠木其它相关文章!

Python 相关文章推荐
CentOS 7下安装Python 3.5并与Python2.7兼容并存详解
Jul 07 Python
pandas的唯一值、值计数以及成员资格的示例
Jul 25 Python
Python中的枚举类型示例介绍
Jan 09 Python
Django csrf 两种方法设置form的实例
Feb 03 Python
python下载微信公众号相关文章
Feb 26 Python
Python一个简单的通信程序(客户端 服务器)
Mar 06 Python
python使用adbapi实现MySQL数据库的异步存储
Mar 19 Python
解决Python设置函数调用超时,进程卡住的问题
Aug 08 Python
详解python3中用HTMLTestRunner.py报ImportError: No module named 'StringIO'如何解决
Aug 27 Python
Django实现基于类的分页功能
Oct 31 Python
python opencv如何实现图片绘制
Jan 19 Python
Python中SQLite如何使用
May 27 Python
python实现经纬度采样的示例代码
Dec 10 #Python
Python urlopen()参数代码示例解析
Dec 10 #Python
Python urllib request模块发送请求实现过程解析
Dec 10 #Python
Python APScheduler执行使用方法详解
Dec 10 #Python
flask项目集成swagger的方法
Dec 09 #Python
python的dict判断key是否存在的方法
Dec 09 #Python
python RSA加密的示例
Dec 09 #Python
You might like
php数组函数序列之array_splice() - 在数组任意位置插入元素
2011/11/07 PHP
批量去除PHP文件中bom的PHP代码
2012/03/13 PHP
PHP中的gzcompress、gzdeflate、gzencode函数详解
2014/07/29 PHP
PHP技术开发微信公众平台
2015/07/22 PHP
PHP递归遍历指定文件夹内的文件实现方法
2016/11/15 PHP
浅谈Laravel队列实现原理解决问题记录
2017/08/19 PHP
php使用curl伪造浏览器访问操作示例
2019/09/30 PHP
Js 中debug方式
2010/02/07 Javascript
基于jQuery.Hz2Py.js插件实现的汉字转拼音特效
2015/05/07 Javascript
JavaScript使用encodeURI()和decodeURI()获取字符串值的方法
2015/08/04 Javascript
纯HTML5制作围住神经猫游戏-附源码下载
2015/08/23 Javascript
学习JavaScript设计模式之中介者模式
2016/01/14 Javascript
AngularJS 实现JavaScript 动画效果详解
2016/09/08 Javascript
有关文件上传 非ajax提交 得到后台数据问题
2016/10/12 Javascript
jQueryeasyui 中如何使用datetimebox 取两个日期间相隔的天数
2017/06/13 jQuery
View.post() 不靠谱的地方你知道多少
2017/08/29 Javascript
详解webpack之图片引入-增强的file-loader:url-loader
2018/10/08 Javascript
详解原生JS回到顶部
2019/03/25 Javascript
python 3.7.0 下pillow安装方法
2018/08/27 Python
浅谈Scrapy网络爬虫框架的工作原理和数据采集
2019/02/07 Python
pytorch中的transforms模块实例详解
2019/12/31 Python
tensorflow使用range_input_producer多线程读取数据实例
2020/01/20 Python
python GUI库图形界面开发之PyQt5简单绘图板实例与代码分析
2020/03/08 Python
css3的@media属性实现页面响应式布局示例代码
2014/02/10 HTML / CSS
仿CSDN Blog返回页面顶部功能实现原理及代码
2013/06/30 HTML / CSS
学前教育毕业生自荐信
2013/10/29 职场文书
高级人员简历的自我评价分享
2013/11/03 职场文书
遗体告别仪式答谢词
2014/01/23 职场文书
生产厂长岗位职责
2014/02/21 职场文书
产品销售计划书
2014/05/04 职场文书
企业宣传标语
2014/06/09 职场文书
四风批评与自我批评范文
2014/10/14 职场文书
北京颐和园导游词
2015/01/30 职场文书
2015年幼儿园安全工作总结
2015/05/12 职场文书
植物园观后感
2015/06/11 职场文书
学术研讨会主持词
2015/07/04 职场文书