显存优化 | Pytorch的显存机制torch.cuda.empty_cache及周边概念
注：文中涉及一些内部底层实现机制，可能和大家外界看到的不一样，这里略过不做介绍。借着笔记，分享平时碰到的技术点，不高端，不炫酷，对你有用更好了。
最近在做模型的优化工作，主要涉及精度，速度和显存优化，今天主要针对显存优化分享下关于Pytorch显存机制的笔记。
显存优化还涉及显存量查看，训练中显存使用分布查看，还有一些节省显存的tricks。我在这篇文章中没有体现，mentor看到文章“嫌弃”说“看的不过瘾”。大佬的视角果然要求高♀️。
我只能弱弱回复：慢慢来。下次碰到其他的再分享也不迟。第一次写，瑟瑟发抖。
0x00 问题
大模型训练容易爆显存，模型训练batch size设置受限，需要短时间内解决眼前的爆炸问题。 （模型本身及数据输入都比较大）
0001 现象
bs较小时候，训练推理正常；bs增大调试过程中，当bs为某个值时，训练正常，但推理会显存爆炸；再增大bs时候，直接爆炸。
注：文中的推理，也可理解为验证或测试。
0010 猜测
- 模型某层申请了特别大的 tensor，可以考虑释放它的缓存。
- 训练过程占用了较多缓存，可以考虑推理前释放训练的缓存。
0011 方法
#释放未被使用的缓存，tensor占用的显存不会被释放
#可以减少碎片，提高利用率
torch.cuda.empty_cache()
注：这只是一个解决的切入口，速效丸，并不是优化显存的最优/最根本方案。
0100 结果
在训练结束，推理前，加上：
torch.cuda.empty_cache()
释放未被使用的缓存，重新分配。一定程度解决增大bs时显存爆炸问题，满足暂时需求；
更进一步增大bs的训练需要在模型本身和更底层分配机制入手。
0101 注意
非必要情况不要经常用，比如可以训练结束调用一下。
该操作会有设备同步，某些场景下会增加时间消耗，具体问题具体分析。
0110 理解
Pytorch可以自动回收我们不用的显存，类似于Python的引用机制，当某一内存内的数据不再有任何变量引用时，这部分的内存便会被释放[1]。
但有一点需要注意，当有一部分显存不再使用的时候，这部分释放的显存通过Nvidia-smi命令是看不到的，举个：
device = torch.device('cuda:0')
# 定义两个tensor
# 120*3*512*512*4/1000/1000 = 377.48M
tensor_4 = torch.randn(120, 3, 512, 512).float().to(device)
# 80*3*512*512*4/1000/1000 = 251.64M
tensor_2 = torch.randn(80, 3, 512, 512).float().to(device)
# 释放
tensor_4 = tensor_4.cpu()
tensor_2 = tensor_2.cpu()
# 这里虽然将上面的显存释放了，但是我们通过Nvidia-smi命令看到显存依然在占用
torch.cuda.empty_cache()
# 只有执行完上面这句，显存才会在Nvidia-smi中释放
Pytorch的开发者也对此进行说明了，这部分释放后的显存通常可以用，只不过不在Nvidia-smi中显示。
也可以理解为释放的显存没有还给Nvidia，被Pytorch当做缓存使用。
但是缓存会有碎片，比如10个1M的tensor释放了不一定能凑成一个10M的空间。还可能有底层分配机制的原因，所以有时候分配不出来，导致空间不够用。
因此需要手动清空下，缓存清空以后就可以从Nvidia再分配。
0111 显卡和显存、GPU的关系[2]
显卡主要由PCB板、GPU、显存构成。
接触过硬件专业的同学应该比较熟悉PCB板，焊板子嘛，焊锡，电烙铁，松香…
GPU是一个附属型的处理器，主要处理计算机中与图形计算有关的工作，并将数据更好地呈现在显示器中。
在处理信息的过程中，它会产生大量的临时数据（未处的、正在处理的、已经处理完成的），这就需要一个专门的地方来存放这些临时数据，即显存。
显卡可能是一个芯片，也可能只是芯片的一部分，具体看硬件的设计，有独立显卡或者集成显卡。
1000 显存和GPU的关系[3]
显存和GPU的关系类似于内存和CPU的关系。
显存用来存放模型、数据，显存越大，所能够运行的网络就越大。
神经网络的显存占用包括但不限于：
1）模型参数的显存占用：只有有参数的层才会有显存占用，这部分的显存占用和输入无关，模型加载完之后就会占用。
有参数的层：卷积层，全连接层，BatchNorm, Embedding层等。
无参数的层：激活层sigmoid（sigmoid，relu），池化层，Dropout等
2）梯度与动量。
3）输入输出。
1001 显存和缓存的关系
GPU计算会频繁使用显存，从显卡中申请一部分显存作为缓存，用完后释放。
如果GPU每次计算都从设备分配显存，那么计算的开销较大，运行速度较慢。如果一开始分配出所有的显存，那么显存占用很高，而利用率很低。
因此有了缓存机制。Pytroch和tensorflow都有自己的缓存机制。
Tensorflow的缓存机制BFC，具体可参考：
https://blog.csdn.net/qq_33096883/article/details/77479647
关于Pytorch 的内存管理机制暂时没有看到具体的定义，后续看到再补充，现在主要涉及就是torch.cuda.empty_cache() 的使用。
附：Pytorch其他节省显存的小技巧，后续有机会实验后再来深究：
https://www.v2ex.com/amp/t/532315
参考链接：
[1]https://blog.csdn.net/qq_29007291/article/details/90451890
[2]http://ask.zol.com.cn/x/5894483.html
[3]https://blog.csdn.net/XYKenny/article/details/89003363?depth_1-utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromBaidu-1&utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromBaidu-1