弄清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的发展史
Sep 26 Python
Python自动化测试Eclipse+Pydev 搭建开发环境
Aug 15 Python
使用Python AIML搭建聊天机器人的方法示例
Jul 09 Python
python:接口间数据传递与调用方法
Dec 17 Python
Python字符串逆序的实现方法【一题多解】
Feb 18 Python
python语言基本语句用法总结
Jun 11 Python
Python 使用 PyMysql、DBUtils 创建连接池提升性能
Aug 14 Python
利用rest framework搭建Django API过程解析
Aug 31 Python
pytorch 图像预处理之减去均值,除以方差的实例
Jan 02 Python
使用Keras训练好的.h5模型来测试一个实例
Jul 06 Python
pytorch中index_select()的用法详解
Jan 06 Python
在Windows下安装配置CPU版的PyTorch的方法
Apr 02 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扩展开发经验分享
2012/09/06 PHP
PHP超级全局变量数组小结
2012/10/04 PHP
Linux系统下PHP-FPM的安装和配置教程
2015/08/17 PHP
php 中的closure用法详解
2017/06/12 PHP
详解cookie验证的php应用的一种SSO解决办法
2017/10/20 PHP
为jquery.ui.dialog 增加“在当前鼠标位置打开”的功能
2009/11/24 Javascript
指定位置如果有图片显示图片,无图片显示广告的JS
2010/06/05 Javascript
JavaScript实现表格排序方法
2013/06/14 Javascript
利用jquery包将字符串生成二维码图片
2013/09/12 Javascript
jquery ajax post提交数据乱码
2013/11/05 Javascript
详谈jQuery操纵DOM元素属性 attr()和removeAtrr()方法
2015/01/22 Javascript
几种经典排序算法的JS实现方法
2016/03/25 Javascript
jquery实现图片上传前本地预览功能
2016/05/10 Javascript
微信小程序 本地数据读取实例
2017/04/27 Javascript
浅谈Angular4中常用管道
2017/09/27 Javascript
JS实现点击链接切换显示隐藏内容的方法
2017/10/19 Javascript
javascript函数的节流[throttle]与防抖[debounce]
2017/11/15 Javascript
vue实例中data使用return包裹的方法
2018/08/27 Javascript
微信小程序如何刷新当前界面的实现方法
2019/06/07 Javascript
原生JavaScript实现留言板
2021/01/10 Javascript
python计算方程式根的方法
2015/05/07 Python
Python实现ssh批量登录并执行命令
2016/10/25 Python
python爬虫爬取淘宝商品信息
2018/02/23 Python
Python读取Word(.docx)正文信息的方法
2018/03/15 Python
python中map的基本用法示例
2018/09/10 Python
Python实现socket非阻塞通讯功能示例
2019/11/06 Python
Python3并发写文件与Python对比
2019/11/20 Python
tensorflow使用指定gpu的方法
2020/02/04 Python
Python网络爬虫信息提取mooc代码实例
2020/03/06 Python
Python多进程编程常用方法解析
2020/03/26 Python
python+selenium 简易地疫情信息自动打卡签到功能的实现代码
2020/08/22 Python
环保倡议书500字
2014/05/15 职场文书
2014年社区计生工作总结
2014/11/18 职场文书
大学辅导员述职报告
2015/01/10 职场文书
如何使用JavaScript策略模式校验表单
2021/04/29 Javascript
Mysql忘记密码解决方法
2022/02/12 MySQL