PyTorch 节省显存
原始文档:https://www.yuque.com/lart/ugkv9f/nvffyf 声明: 大部分内容来自知乎和其他博客的分享, 这里只作为一个收集罗列. 欢迎给出更多建议.
知乎回答(欢迎点赞哦):
- pytorch dataloader 数据加载占用了大部分时间, 各位大佬都是怎么解决的? - 人民艺术家的回答 - 知乎
使用 pytorch 时, 训练集数据太多达到上千万张, Dataloader 加载很慢怎么办? - 人民艺术家的回答 - 知乎
更新日志
2019 年 11 月 29 日: 更新一些模型设计技巧和推理加速的内容, 补充了下 apex 的一个介绍链接,
另外删了 tfrecord, pytorch 能用么? 这个我记得是不能, 所以删掉了(表示删掉:<)- 2019 年 11 月 30 日: 补充 MAC 的含义, 补充 ShuffleNetV2 的论文链接
- 2019 年 12 月 02 日: 之前说的 pytorch 不能用 tfrecord, 今天看到https://www.zhihu.com/question/358632497下的一个回答, 涨姿势了
- 2019 年 12 月 23 日: 补充几篇关于模型压缩量化的科普性文章
- 2020 年 2 月 7 日: 从文章中摘录了一点注意事项, 补充在了代码层面小节
- 2020 年 4 月 30 日:
- 添加了一个 github 的文档备份
- 补充了卷积层和 BN 层融合的介绍的链接
- 另外这里说明下, 对于之前参考的很多朋友的文章和回答, 没有把链接和对应的内容提要关联在一起, 估计会导致一些朋友阅读时相关的内容时的提问, 无法问到原作者, 这里深感抱歉.
- 调整部分内容, 将内容尽量与参考链接相对应
- 2020 年 5 月 18 日: 补充一些关于 PyTorch 节省显存的技巧. 同时简单调整格式. 另外发现一个之前的错误:
non_blocking=False的建议应该是non_blocking=True. - 2021 年 01 月 06 日：调整下关于读取图片数据的一些介绍.
- 2021 年 01 月 13 日：补充了一条推理加速的策略. 我觉得我应该先更新 github 的文档，知乎答案的更新有点麻烦，也没法比较更改信息，就很费劲。
- 2022 年 6 月 26 日：重新调整了下格式和内容安排，同时补充了更多的参考资料和一些最新发现的有效内容。
2022年7月3日：重排格式，调整了重参数化和管理显存的部分。
代码层面
使用In-Place操作
对于默认支持
inplace的操作尽量启用。比如relu可以使用inplace=True。可以将
batchnorm和一些特定的激活函数打包成[inplace_abn](https://github.com/mapillary/inplace_abn)。损失函数
每次循环结束时删除 loss, 可以节约很少显存, 但聊胜于无。可见Tensor to Variable and memory freeing best practices
混合精度
可以节约一定的显存并提速, 但是要小心一些不安全的操作如 mean 和 sum。
混合精度训练的介绍文章：
[NVIDIA/Apex](https://github.com/nvidia/apex)提供的混合精度支持。PyTorch1.6 开始提供的
[torch.cuda.amp](https://pytorch.org/docs/stable/notes/amp_examples.html)以支持混合精度。优化梯度计算
基本优化
对于不需要反向传播的前向阶段，如验证和推理期间，使用
torch.no_grad来包裹代码。- 注意
model.eval()不等于torch.no_grad(), 请看如下讨论: ‘model.eval()’ vs ‘with torch.no_grad()’
- 注意
将不需要计算梯度的变量的
requires_grad设为False, 让变量不参与梯度的后向传播，以减少不必要的梯度的显存占用。移除不必要的梯度路径
典型工作为SBP，即Stochastic Backpropagation: A Memory Efficient Strategy for Training Video Models。SBP保持所有前向传播路径，但在每个训练步骤中随机和独立地删除每个网络层的反向传播路径。它通过消除对应于反向传播路径的缓存激活值的需求，从而降低GPU存储成本。而这个量可以通过可调节的保持率来控制。
更详细的解读可见：Stochastic Backpropagation: A Memory Efficient Strategy for Training Video Models
Model Trick | CVPR 2022 Oral - Stochastic Backpropagation A Memory Efficient Strategy
自定义包含更多细粒度操作的Function
这部分内容主要在Vision Longformer中有所讨论。在Vision Longformer中的实现中，作者们对比了针对提出的结构不同的实现方式随着序列长度的运行时间和显存占用的情况。CUDA实现显存占用是最有优势但运行时间却并非最优。最优的反而是利用PyTorch Python端自定义算子的方式，即继承
Function并自定义前向传播和反向传播。
时间上优势主要是因为PyTorch自带算子已经被高度优化，因而可以获得更好的时间效益。另外这种形式也可以通过手动实现反向传播来进一步节省显存。这是因为前向传播中autograd是按照细粒度算子的层级来记录中间结果的，实际上如果按照整个模块的角度来考虑的话，有时只需存储部分结果用于反向传播即可，所以手动实现更具空间优势。显存清理
可以使用
del删除不必要的中间变量。- How can we release GPU memory cache?：
torch.cuda.empty_cache()这是del的进阶版，使用nvidia-smi会发现显存有明显的变化。但是训练时最大的显存占用似乎没变。 PyTorch显存占用分析，PyTorch显存机制分析：PyTorch在进行深度学习训练的时候，有4大部分的显存开销，分别是模型参数(parameters)，模型参数的梯度(gradients)，优化器状态(optimizer states)以及中间激活值(intermediate activations) 或者叫中间结果(intermediate results)。
梯度累加（Gradient Accumulation）
把一个
batchsize=64分为两个 32 的 batch，两次 forward 以后，backward 一次。但会影响batchnorm等和batchsize相关的层。
在PyTorch 的文档中提到了梯度累加与混合精度并用的例子。
使用梯度累加技术可以对分布式训练加速，这可以参考：[原创][深度][PyTorch] DDP 系列第三篇：实战与技巧 - 996 黄金一代的文章 - 知乎梯度检查点（Gradient Checkpointing）
PyTorch 中提供了
[torch.utils.checkpoint](https://pytorch.org/docs/1.11/checkpoint.html#torch-utils-checkpoint)。这是通过在反向传播期间，在每个检查点位置重新执行一次前向传播来实现的。
论文Training Deep Nets with Sublinear Memory Cost基于梯度检查点技术，将显存从 O(N) 降到了 O(sqrt(N))。对于越深的模型, 这个方法省的显存就越多, 且速度不会明显变慢。- torch.utils.checkpoint 简介 和 简易使用
Sublinear Memory Cost 的一份 PyTorch 实现，参考自：Pytorch 有什么节省内存(显存)的小技巧? - Lyken 的回答 - 知乎
相关工具
These codes can help you to detect your GPU memory during training with Pytorch. https://github.com/Oldpan/Pytorch-Memory-Utils
Just less than nvidia-smi? https://github.com/wookayin/gpustat
参考资料
- 浅谈深度学习: 如何计算模型以及中间变量的显存占用大小
- 如何在 Pytorch 中精细化利用显存
- Pytorch 有什么节省显存的小技巧? - 陈瀚可的回答 - 知乎
- PyTorch 显存机制分析 - Connolly 的文章 - 知乎
若有收获，就点个赞吧
0 人点赞
