弄清Pytorch显存的分配机制
更新时间：2020年12月10日 10:18:22 作者：颀周
这篇文章主要介绍了Pytorch显存的分配机制的相关资料，帮助大家更好的理解和使用Pytorch，感兴趣的朋友可以了解下
对于显存不充足的炼丹研究者来说，弄清楚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())
输出如下：
代码首先分配3GB的显存创建变量x，然后计算y，再用y进行反向传播。可以看到，创建x后与计算y后分别占显存3GB与6GB，这是合理的。另外，后面通过backward()，计算出x.grad，占存与x一致，所以最终一共占有显存9GB，这也是合理的。但是，输出显示了显存的峰值为12GB，这多出的3GB是怎么来的呢？首先画出计算图：
下面通过列表的形式来模拟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 显存分配的资料请关注脚本之家其它相关文章！
相关文章
这篇文章主要为大家介绍了Python运算符，具有一定的参考价值，感兴趣的小伙伴们可以参考一下，希望能够给你带来帮助2021-12-12
在本篇文章里，小编给大家分享的是一篇关于python爬虫是否容易学的相关知识点内容，有兴趣的朋友们可以阅读下。2020-06-06
这篇文章主要介绍了Python复制Word内容并使用格式设字体与大小实例代码，小编觉得还是挺不错的，具有一定借鉴价值，需要的朋友可以参考下2018-01-01
在本篇文章里小编给大家整理的是关于Python实现新型冠状病毒传播模型及预测代码内容，有兴趣的朋友们可以学习下。2020-02-02
今天小编就为大家分享一篇Python创建一个元素都为0的列表实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧2019-11-11
python读取和保存为excel、csv、txt文件及对DataFrame文件的基本操作指南
最近在做一个项目,必须把结果保存到excel文件中,下面这篇文章主要给大家介绍了关于python读取和保存为excel、csv、txt文件及对DataFrame文件的基本操作指南的相关资料,文中通过实例代码介绍的非常详细,需要的朋友可以参考下2022-08-08
这篇文章主要介绍了如何利用Python编写一个藏头诗在线生成器，文中的示例代码讲解详细，感兴趣的同学可以跟随小编一起动手尝试一下2022-04-04
这篇文章主要介绍了Python json格式化打印实现过程解析,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下2020-07-07
本文给大家介绍的是利用urllib模块通过指定的URL抓取网页内容 所谓网页抓取,就是把URL地址中指定的网络资源从网络流中读取出来,保存到本地，有需要的小伙伴可以参考下2016-04-04
这篇文章主要介绍了Anaconda超详细保姆级安装配置教程,本篇教程对于新手很友好强烈推荐,看完这篇文章相信你对如何安装配置Anaconda一定了然于胸,我们一起来看看吧2023-03-03
最新评论