zoukankan      html  css  js  c++  java
  • 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:

  • 相关阅读:
    Servlet学习笔记(1)--第一个servlet&&三种状态对象(cookie,session,application)&&Servlet的生命周期
    XML学习笔记(2)--dom4j操作XML
    坦克大战(版本2.5-版本2.9)
    坦克大战(版本1.7-版本2.4)
    坦克大战(版本1.0-版本1.6)
    坦克大战(版本0.1-版本0.9)
    JavaSE聊天室
    HTML+CSS+JS学习总结
    JDBC学习笔记(10)——调用函数&存储过程
    VS工程的相对路径写法
  • 原文地址:https://www.cnblogs.com/qizhou/p/14110086.html
Copyright © 2011-2022 走看看