弄清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多线程编程(一):threading模块综述
Apr 05 Python
使用Python编写简单的端口扫描器的实例分享
Dec 18 Python
python 创建弹出式菜单的实现代码
Jul 11 Python
python 随机数使用方法,推导以及字符串,双色球小程序实例
Sep 12 Python
Python聊天室程序(基础版)
Apr 01 Python
Python 分发包中添加额外文件的方法
Aug 16 Python
Python2与Python3的区别点整理
Dec 12 Python
Python 2种方法求某个范围内的所有素数(质数)
Jan 31 Python
利用Python自动化操作AutoCAD的实现
Apr 01 Python
python中的yield from语法快速学习
Nov 06 Python
python如何发送带有附件、正文为HTML的邮件
Feb 27 Python
Pytest之测试命名规则的使用
Apr 16 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
YII Framework框架教程之缓存用法详解
2016/03/14 PHP
php 判断IP为有效IP地址的方法
2018/01/28 PHP
JavaScript中的History历史对象
2008/01/16 Javascript
用cssText批量修改样式
2009/08/29 Javascript
基于jquery的合并table相同单元格的插件(精简版)
2011/04/05 Javascript
jQuery$命名冲突怎么办如何解决
2014/01/16 Javascript
window.close(); 关闭浏览器窗口js代码的总结介绍
2016/07/14 Javascript
BootStrap的table表头固定tbody滚动的实例代码
2016/08/24 Javascript
常用的javascript设计模式
2017/01/11 Javascript
vue2.0获取自定义属性的值
2017/03/28 Javascript
jQuery查找和过滤_动力节点节点Java学院整理
2017/07/04 jQuery
angular2系列之路由转场动画的示例代码
2017/11/09 Javascript
vue表单验证你真的会了吗?vue表单验证(form)validate
2019/04/07 Javascript
详解jquery和vue对比
2019/04/16 jQuery
通过说明与示例了解js五种设计模式
2019/06/17 Javascript
js DOM的事件常见操作实例详解
2019/12/16 Javascript
[36:33]2018DOTA2亚洲邀请赛 4.3 突围赛 EG vs Newbee 第二场
2018/04/04 DOTA
用Python编写分析Python程序性能的工具的教程
2015/04/01 Python
python模块之StringIO使用示例
2015/04/08 Python
PyQt5每天必学之切换按钮
2020/08/20 Python
pandas将numpy数组写入到csv的实例
2018/07/04 Python
500行代码使用python写个微信小游戏飞机大战游戏
2019/10/16 Python
python输出pdf文档的实例
2020/02/13 Python
Django Model中字段(field)的各种选项说明
2020/05/19 Python
python tkiner实现 一个小小的图片翻页功能的示例代码
2020/06/24 Python
python基于opencv 实现图像时钟
2021/01/04 Python
迷你唐卡软皮鞋:Minnetonka Moccasin
2018/05/01 全球购物
What is the purpose of Void class? Void类的作用是什么?
2016/10/31 面试题
linux面试题参考答案(8)
2015/08/11 面试题
什么是组件架构
2016/05/15 面试题
酒店执行总经理岗位职责
2013/12/15 职场文书
司马光教学反思
2014/02/01 职场文书
电视购物广告词
2014/03/19 职场文书
师德师风主题教育活动总结
2015/05/07 职场文书
培训讲师开场白
2015/06/01 职场文书
岁月神偷观后感
2015/06/11 职场文书