弄清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实现的生成自我描述脚本分享(很有意思的程序)
Jul 18 Python
使用Python的Flask框架表单插件Flask-WTF实现Web登录验证
Jul 12 Python
python实现实时监控文件的方法
Aug 26 Python
python 中if else 语句的作用及示例代码
Mar 05 Python
tensorflow 1.0用CNN进行图像分类
Apr 15 Python
对Pandas MultiIndex(多重索引)详解
Nov 16 Python
简单了解Django应用app及分布式路由
Jul 24 Python
win10环境下配置vscode python开发环境的教程详解
Oct 16 Python
Python搭建代理IP池实现获取IP的方法
Oct 27 Python
python为Django项目上的每个应用程序创建不同的自定义404页面(最佳答案)
Mar 09 Python
python json.dumps() json.dump()的区别详解
Jul 14 Python
Python基础数据类型tuple元组的概念与用法
Aug 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在Web开发领域的优势
2006/10/09 PHP
PHP6 先修班 JSON实例代码
2008/08/23 PHP
php类常量的使用详解
2013/06/08 PHP
Laravel 5框架学习之表单验证
2015/04/08 PHP
PHP内存使用情况如何获取
2015/10/10 PHP
PHP二维数组去重算法
2016/12/17 PHP
学习thinkphp5.0验证类使用方法
2017/11/16 PHP
PHP中有关长整数的一些操作教程
2019/09/11 PHP
对字符串进行HTML编码和解码的JavaScript函数
2010/02/01 Javascript
jQuery实现平滑滚动到指定锚点的方法
2015/03/20 Javascript
jQuery地图map悬停显示省市代码分享
2015/08/20 Javascript
vue2.0多条件搜索组件使用详解
2020/03/26 Javascript
浅析node Async异步处理模块用例分析及常用方法介绍
2017/11/17 Javascript
原生js实现拖拽功能基本思路详解
2018/04/18 Javascript
详解vue2.0监听属性的使用心得及搭配计算属性的使用
2018/07/18 Javascript
对Vue2 自定义全局指令Vue.directive和指令的生命周期介绍
2018/08/30 Javascript
vue学习笔记之slot插槽用法实例分析
2020/02/29 Javascript
vue实践---根据不同环境,自动转换请求的url地址操作
2020/09/21 Javascript
python中global与nonlocal比较
2014/11/21 Python
在java中如何定义一个抽象属性示例详解
2017/08/18 Python
python实现简单点对点(p2p)聊天
2017/09/13 Python
在django admin中添加自定义视图的例子
2019/07/26 Python
Python 二叉树的层序建立与三种遍历实现详解
2019/07/29 Python
python使用yield压平嵌套字典的超简单方法
2019/11/02 Python
python判断两个序列的成员是否一样的实例代码
2020/03/01 Python
在Ubuntu中安装并配置Pycharm教程的实现方法
2021/01/06 Python
官方授权图形T恤和服装:Fifth Sun
2019/06/12 全球购物
英国玛莎百货澳大利亚:Marks & Spencer Australia
2019/08/30 全球购物
be2台湾单身男女交友:全球网路婚姻介绍的领导品牌
2019/10/11 全球购物
经典优秀个人求职自荐信格式
2013/09/25 职场文书
应届生法律顾问求职信
2013/11/19 职场文书
幼儿运动会邀请函
2014/01/17 职场文书
机械专业应届毕业生自荐书
2014/06/12 职场文书
经典禁毒标语
2014/06/16 职场文书
离婚协议书怎样才有法律效力
2014/10/10 职场文书
鸡毛信观后感
2015/06/11 职场文书