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：