Pytorch节省显存 - checkpoint

Pytorch节省显存 - checkpoint

官方文档: torch.utils.checkpoint.checkpoint

一、查看模型及中间变量的显存占用情况

1. 显存占用

参考博客:
https://oldpan.me/archives/how-to-calculate-gpu-memory
https://zhuanlan.zhihu.com/p/31558973

主要由两个部分组成：

• 模型参数
• 模型计算产生的中间变量：
  • 前向计算过程中，隐层的输出都被保存在显存中。
  • 后向梯度计算时，需要用到这些隐层的输出。梯度和优化器动量，将占据一部分显存。

显存占用可简化为：

显存占用 = 模型参数 + batch_size * 模型计算产生的中间变量

2. 训练阶段显存为何不足

参考博客:
https://www.yanxishe.com/columnDetail/18777

举例说明：一次训练过程中的显存使用情况

• 前向计算时，大量的隐层输出在显存中累积，导致显存很快被占满。
• 后向计算时，隐层输出被用来计算梯度。此时，不再被用到的隐层输出和梯度，会被释放显存。
  pytorch自动回收不用的显存。当某一内存内的数据不再有任何变量引用时，这部分的内存便会被释放。因此，中间变量逐渐减少。

3. 显存跟踪代码

参考博客: https://oldpan.me/archives/pytorch-gpu-memory-usage-track
项目地址: https://github.com/Oldpan/Pytorch-Memory-Utils

理解

https://www.yanxishe.com/columnDetail/18777

使用举例

简单使用

多参数使用

返回值要求

在DPP中的使用

checkpoint find_unused_parameters=False
RuntimeError: Expected to mark a variable ready only once. This error is caused by use of a module parameter outside the forward function.
https://github.com/allenai/longformer/issues/63

TypeError: CheckpointFunctionBackward.forward: expected Variable (got dict) for return value 0

减少显存
https://github.com/open-mmlab/mmdetection/blob/927d71a98ce2446193c626d26268df0c29922fb0/mmdet/core/fp16/hooks.py#L11

加载模型
https://www.zhihu.com/question/351342218

https://www.zhihu.com/question/274635237
https://zhuanlan.zhihu.com/p/232466530

https://blog.csdn.net/qq_28660035/article/details/80688427

追踪显存
https://github.com/Oldpan/Pytorch-Memory-Utils
https://oldpan.me/archives/pytorch-gpu-memory-usage-track

前向推理
http://noahsnail.com/2020/05/29/2020-05-29-Pytorch%E4%B8%ADrequires_grad_(),%20detach(),%20torch.no_grad()%E7%9A%84%E5%8C%BA%E5%88%AB/

显存释放
https://www.cnblogs.com/jiangkejie/p/11430673.html

apex
https://github.com/NVIDIA/apex

https://fyubang.com/2019/08/26/fp16/

https://gitlab.hk.sensetime.com/denghanming/PointCloudDet3D/commit/5cf05f7c09bb0eaa686768066ae24600a9dc58c3

python fault handler