弄清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 相关文章推荐
python实现socket端口重定向示例
Feb 10 Python
在Mac OS上搭建Python的开发环境
Dec 24 Python
python解决Fedora解压zip时中文乱码的方法
Sep 18 Python
Python绘制并保存指定大小图像的方法
Jan 10 Python
Python字典循环添加一键多值的用法实例
Jan 20 Python
django框架事务处理小结【ORM 事务及raw sql,customize sql 事务处理】
Jun 27 Python
python线程信号量semaphore使用解析
Nov 30 Python
Python 音频生成器的实现示例
Dec 24 Python
Python matplotlib实时画图案例
Apr 23 Python
python用opencv完成图像分割并进行目标物的提取
May 25 Python
Python调用jar包方法实现过程解析
Aug 11 Python
Python采集壁纸并实现炫轮播
Apr 30 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中使用exec,system等函数调用系统命令的方法(不建议使用,可导致安全问题)
2012/09/07 PHP
PHP获取当前时间不准确问题解决方案
2020/08/14 PHP
javascript入门·动态的时钟,显示完整的一些方法,新年倒计时
2007/10/01 Javascript
jsTree 基于JQuery的排序节点 Bug
2011/07/26 Javascript
JS中的prototype与面向对象的实例讲解
2013/05/22 Javascript
jQuery获取和设置表单元素的方法
2014/02/14 Javascript
Jquery性能优化详解
2014/05/15 Javascript
jquery实现表单输入时提示文字滑动向上效果
2015/08/10 Javascript
JavaScript中的操作符类型转换示例总结
2016/05/30 Javascript
原生javascript实现文件异步上传的实例讲解
2017/10/26 Javascript
Vue中render方法的使用详解
2018/01/26 Javascript
vue.js 实现评价五角星组件的实例代码
2018/08/13 Javascript
js中this的指向问题归纳总结
2018/11/28 Javascript
详解vue2.6插槽更新v-slot用法总结
2019/03/09 Javascript
详解微信小程序scroll-view横向滚动的实践踩坑及隐藏其滚动条的实现
2019/03/14 Javascript
ES6实现图片切换特效代码
2020/01/14 Javascript
javaScript实现一个队列的方法
2020/07/14 Javascript
微信小程序自定义底部弹出框动画
2020/11/18 Javascript
在Python的Flask框架中实现全文搜索功能
2015/04/20 Python
python 环境变量和import模块导入方法(详解)
2017/07/11 Python
python对配置文件.ini进行增删改查操作的方法示例
2017/07/28 Python
python的文件操作方法汇总
2017/11/10 Python
解决pycharm每次新建项目都要重新安装一些第三方库的问题
2019/01/17 Python
python 解决mysql where in 对列表(list,,array)问题
2020/06/06 Python
python 19个值得学习的编程技巧
2020/08/15 Python
浅谈Python __init__.py的作用
2020/10/28 Python
Myprotein蛋白粉美国官网:欧洲畅销运动营养品牌
2016/11/15 全球购物
应届生煤化工求职信
2013/10/21 职场文书
公司授权委托书范文
2014/08/02 职场文书
查摆剖析材料范文
2014/09/30 职场文书
经典搞笑版检讨书
2015/02/19 职场文书
2015年办公室人员工作总结
2015/05/15 职场文书
html5实现点击弹出图片功能
2021/07/16 HTML / CSS
Node-Red实现MySQL数据库连接的方法
2021/08/07 MySQL
css3属性选择器 “~”(波浪号) “,”(逗号) “+”(加号)和 “>”(大于号)
2022/04/19 HTML / CSS
win10忘记pin密码登录不了怎么办?win10忘记pin密码登不进去的解决方法
2022/07/07 数码科技