弄清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实例分享:快速查找出被挂马的文件
Jun 08 Python
Python、Javascript中的闭包比较
Feb 04 Python
Python中的高级函数map/reduce使用实例
Apr 13 Python
python发送邮件实例分享
Jul 28 Python
python的paramiko模块实现远程控制和传输示例
Oct 13 Python
Python使用itertools模块实现排列组合功能示例
Jul 02 Python
Python利用itchat库向好友或者公众号发消息的实例
Feb 21 Python
django多对多表的创建,级联删除及手动创建第三张表
Jul 25 Python
Python测试线程应用程序过程解析
Dec 31 Python
python opencv如何实现图片绘制
Jan 19 Python
基于python计算滚动方差(标准差)talib和pd.rolling函数差异详解
Jun 08 Python
基于Python-turtle库绘制路飞的草帽骷髅旗、美国队长的盾牌、高达的源码
Feb 18 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
PIGCMS 如何关闭聊天机器人
2015/02/12 PHP
分享10段PHP常用代码
2015/11/11 PHP
PHP中PDO事务处理操作示例
2018/05/02 PHP
层序遍历在ExtJs的TreePanel中的应用
2009/10/16 Javascript
IE6弹出“已终止操作”的解决办法
2010/11/27 Javascript
myFocus slide3D v1.1.0 使用方法与下载
2011/01/12 Javascript
修复IE9&safari 的sort方法
2011/10/21 Javascript
深入理解JSON数据源格式
2014/01/10 Javascript
浏览器窗口大小变化时使用resize事件对框架不起作用的解决方法
2014/05/11 Javascript
基于jquery实现的可编辑下拉框实现代码
2014/08/02 Javascript
jQuery异步上传文件插件ajaxFileUpload详细介绍
2015/05/19 Javascript
Jquery跨域获得Json的简单实例
2016/05/18 Javascript
JS实现点击复选框变更DIV显示状态的示例代码
2017/12/18 Javascript
web前端页面生成exe可执行文件的方法
2018/02/08 Javascript
Vue 中使用vue2-highcharts实现曲线数据展示的方法
2018/03/05 Javascript
node puppeteer(headless chrome)实现网站登录
2018/05/09 Javascript
JavaScript实现的前端AES加密解密功能【基于CryptoJS】
2018/08/28 Javascript
vue中eslintrc.js配置最详细介绍
2018/12/21 Javascript
Element输入框带历史查询记录的实现示例
2019/01/15 Javascript
vue+elementUI组件table实现前端分页功能
2020/11/15 Javascript
理解Proxy及使用Proxy实现vue数据双向绑定操作
2020/07/18 Javascript
Python基础学习之类与实例基本用法与注意事项详解
2019/06/17 Python
python numpy 常用随机数的产生方法的实现
2019/08/21 Python
Python 使用threading+Queue实现线程池示例
2019/12/21 Python
基于Python爬取搜狐证券股票过程解析
2020/11/18 Python
python中openpyxl和xlsxwriter对Excel的操作方法
2021/03/01 Python
阿迪达斯德国官方网站:adidas德国
2017/07/12 全球购物
100%有机精油,美容油:House of Pure Essence
2018/10/30 全球购物
Harrods美国:英国最大的百货公司
2018/11/04 全球购物
澳大利亚领先的亚麻品牌:Bed Threads
2019/12/16 全球购物
华为智利官方商店:Huawei Chile
2020/05/09 全球购物
夜大毕业生自我评价分享
2013/11/10 职场文书
一份创业计划书范文
2014/02/08 职场文书
2014年党建工作汇报材料
2014/10/27 职场文书
五年级下册复习计划
2015/01/19 职场文书
Mysql官方性能测试工具mysqlslap的使用简介
2021/05/21 MySQL