Pytorch显存机制与显存占用(一) —— 理论分析(memory)(训练过程中占用显存最大部分的是activation)(pytorch缓存区/缓存分配器)

PyTorch在进行深度学习训练的时候，有4大部分的显存开销，分别是

      ①模型参数(parameters)

      ②模型参数的梯度(gradients)

      ③优化器状态(optimizer states)

      ④中间激活值(intermediate activations) 或者叫中间结果(intermediate results)。

其中activation占绝对大头，50%以上(有些地方说训练过程中占用显存最大的是计算图，把activation看作计算图的节点的话，也没错)

优化器占的大小只是weight的2倍(Adam),SGD啥的话甚至1倍

它们分别在深度学习的4个步骤中产生

      1. 模型定义：定义了模型的网络结构，产生模型参数；

while(你想训练):

      2. 前向传播：执行模型的前向传播，产生中间激活值；

      3. 后向传播：执行模型的后向传播，产生梯度；

      4. 梯度更新：执行模型参数的更新，第一次执行的时候产生优化器状态。

在模型定义完之后，2~4循环执行。

Pytorch context开销

      只要你把任何东西（无论是多小的tensor）放到GPU显存中，那么你至少会占用1000MiB左右的显存（根据cuda版本，会略有不同）。这部分显存是cuda running时固有配件必须要占掉的显存，是无法避免的，这就是pytorch的context开销。

​

      什么是PyTorch context? 其实官方给他的称呼是CUDA context，就是在第一次执行CUDA操作，也就是使用GPU的时候所需要创建的维护设备间工作的一些相关信息

      验证一下

import torch
temp = torch.tensor([1.0]).cuda()

 

         怎么去减小这个开销...官方也给了一个办法，看看自己有哪些cuda依赖是不需要的，比如cuDNN，然后自己重新编译一遍PyTorch。编译的时候把对应的包的flag给设为false就好了。我是还没有试过，要搭编译的环境太难受了，而且还要经常和库做更新。

Pytorch显存分析(不要用nvidia-smi) / 显存释放机制

      用nvidia-smi或者gpustat来看pytorch程序的显存占用不是很合适的

      因为Pytorch的机制是使用缓存分配器来管理缓存分配的(因为这样速度快), 但是在缓存分配器的机制下, 一个Tensor就算被释放了，进程也不会把空闲出来的显存还给GPU，而是等待下一个Tensor来填入这一片被释放的空间(即只要一个Tensor对象在后续不会再被使用，那么PyTorch就会自动回收该Tensor所占用的显存，并以缓冲区的形式继续占用显存，所以在nvidia-smi/gpustat中看到的显存并没有减少)

      即用nvidia-smi/gpustat看到的其实是pytorch缓存区/缓存分配器的情况

     要是实在看缓冲区不爽的话，也可以用torch.cuda.empty_cache()把它归零，或者加一个环境变量PYTORCH_NO_CUDA_MEMORY_CACHING=1，但是程序速度会变慢哦(试过在一个实验里慢了三倍)

验证一下

import torch

device = torch.device('cuda:0')
# 定义两个tensor
dummy_tensor_4 = torch.randn(120, 3, 512, 512).float().to(device)  # 120*3*512*512*4/1000/1000 = 377.48M
dummy_tensor_5 = torch.randn(80, 3, 512, 512).float().to(device)  # 80*3*512*512*4/1000/1000 = 251.64M

# 然后释放
dummy_tensor_4 = dummy_tensor_4.cpu()
dummy_tensor_5 = dummy_tensor_5.cpu()
# 这里虽然将上面的显存释放了，但是我们通过Nvidia-smi命令看到显存依然在占用
torch.cuda.empty_cache()
# 只有执行完上面这句，显存才会在Nvidia-smi中释放
print()

 

 

 还有的1437就是context开销了

 

但是这样在ipython下是没有作用的，可能是有什么特殊的机制

用

torch.cuda.memory_reserved()

可以查看当前进程所分配的显存缓冲区(tensor所占用的显存大小就包含在这里面)。

而

torch.cuda.empty_cache()

会清理当前显存缓冲区中空闲出来的显存(unused显存)，被tensor占用的显存不会清

torch.cuda.memory_cached()的值 从 实验来看是一直等于torch.cuda.memory_reserved()的值的

torch.cuda.memory_reserved()的值并不就等于nvidia-smi显示的值，nvidia-smi值是reserved_memory和torch context显存之和

而我们实际想看的当前进程中Torch.Tensor所占用的GPU显存，应该用它看

torch.cuda.empty_cache() #使用memory_allocated前先清空一下cache
torch.cuda.memory_allocated()

注意默认是cuda:0的，所以如果模型被放到了cuda:2上，那么应该torch.cuda.memory_allocated("cuda:2"), 否则会得出0的

      有什么好处？进程不需要重新向GPU申请显存了，运行速度会快很多，有什么坏处？他不能准确地给出某一个时间点具体的Tensor占用的显存，nvidia-smi显示的而是已经分配到的显存(显存缓冲区)和context开销之和, 也就是reserved_memory和torch context显存之和。

      这也是令很多人在使用PyTorch时对显存占用感到困惑的罪魁祸首。

torch.cuda is all you need

在分析PyTorch的显存时候，一定要使用torch.cuda里的显存分析函数

用的最多的是

torch.cuda.memory_allocated()
和
torch.cuda.max_memory_allocated()

前者可以精准地反馈当前进程中Torch.Tensor所占用的GPU显存，后者则可以告诉我们到调用函数为止所达到的最大的显存占用字节数。

还有

torch.cuda.memory_summary()

查看显存信息

统计的还是挺全面的

Pytorch底层显存分配机制(按页分配)

      在PyTorch中，显存是按页为单位进行分配的，这可能是CUDA设备的限制。就算我们只想申请4字节的显存，pytorch也会为我们分配512字节或者1024字节的空间。

     即就算我们只想申请4字节的显存，pytorch也会先向CUDA设备申请2MB的显存到自己的cache区中，然后pytorch再为我们分配512字节或者1024字节的空间。这个在使用torch.cuda.memory_allocated()的时候可以看出来512字节；用torch.cuda.memory_cached()可以看出向CUDA申请的2MB。直观点来说，如图所示，PyTorch的显存管理是一个层级结构。

import torch

# 模型初始化
linear1 = torch.nn.Linear(1024,1024, bias=False).cuda() # + 4194304
print(torch.cuda.memory_allocated())
linear2 = torch.nn.Linear(1024, 1, bias=False).cuda() # + 4096
print(torch.cuda.memory_allocated())

# 输入定义
inputs = torch.tensor([[1.0]*1024]*1024).cuda() # shape = (1024,1024) # + 4194304
print(torch.cuda.memory_allocated())

# 前向传播
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304 + 512
print(torch.cuda.memory_allocated())

# 后向传播
loss.backward() # memory - 4194304 + 4194304 + 4096
print(torch.cuda.memory_allocated())

# 再来一次~
loss = sum(linear2(linear1(inputs))) # shape = (1) # memory + 4194304  (512没了，因为loss的ref还在)
print(torch.cuda.memory_allocated())
loss.backward() # memory - 4194304
print(torch.cuda.memory_allocated())

在模型的定义部分，显存占用量约为参数占用量*4 (tensor的类型默认是float32，一个float32占4字节)

所以linear1是1024*1024*4 = 4194304

linear2是1024*4 = 4096

所以第一个输出是4194304，第2个输出是4194304+4096=4198400

输入数据是1024*1024，显存占用是1024*1024*4

所以第三个输出是4198400+1024*1024*4=8392704

在前向传播过程，显存增加等于每一层模型产生的结果的显存之和，且跟batch_size成正比

inputs的shape是1*1024*1024, 经过linear1的结果还是1*1024*1024

所以第四个输出是8392704+1024*1024*4=12587008

经过linear2的结果是1*1024

所以第五个输出的12587008+1024*4=12591104

然后loss存储最后的结果求和，按理说loss只用个4字节就好了，但是在PyTorch中，显存是按页为单位进行分配的。就算我们只想申请4字节的显存，pytorch也会为我们分配512字节

所以第六个输出12591104+512=12591616

反向传播中，后向传播会将模型的中间激活值给消耗并释放掉，并为每一个模型中的参数计算其对应的梯度。在第一次执行的时候，会为模型参数分配对应的用来存储梯度的空间。

反向传播部分没太看懂。。。。

以后再补吧

PyTorch显存机制分析 - 知乎