Pytorch内存优化方法，显著提升模型训练batch_size，减少out of Memory错误发生

主要参考资料：

1. 知乎讨论；
2. pytorch论坛；
3. 官方文档；
4. https://mathpretty.com/11156.html

---

使用in_place操作使得Pytorch的allocator不会记录该部分的原tensor，从而减少显存的消耗。也正是因为如此，如果在网络反向计算梯度的过程中需要用到原tensor，那么就不能够使用in_place操作，否则会使得反向传播报错。

Pytorch的nn.ReLU支持in_place操作，一般原则是，只要ReLU的前一层“不需要”进行梯度计算，那么ReLU就可以使用in_place实现：例如，当ReLU层前是一个BN层时（Pytorch的BN会在反向传播过程中重新计算对应的需求项，因此无需记录梯度信息）。

---

torch.utils.checkpoint提供了一种用 “时间换空间” 的解决思路，即对于部分或者，我们可以用对其进行包装，被包装的model，在正向传播的过程中不会记录对应的结果（节省空间），而是在反向传播的过程中重新计算（增加了计算成本）。
除此以外，使用checkpoint还有很多需要注意的细节，如果使用不当可能造成“两次”计算的结果是不相等的，这样就产生了误差，是我们不想看到的，更多细节可以参考 官方文档 。
小案例：

这里在推荐两个基于Pytorch checkpoint特性实现的giuhub库，实现了诸如将 batchnorm 和relu打包成inplace操作等：
https://github.com/gpleiss/efficient_densenet_pytorch
https://github.com/mapillary/inplace_abn

---

根据大家的说法，这个方法其实作用并不是很大，但是我们姑且还是“了解一下”，其实它的思想很简单，一次循环之后如果我们已经完成了对应的optim以及loss追溯的操作，那么output， loss等值当然就不再被需要了，这时候手动释放这些内存可能是值得尝试的。

注意：如果数据集很小，这个操作可能得不偿失

---

一般而言，Pytorch的模型训练要求数据为float32的精度，但事实上不同的训练子层对数据精度的需求是不一致的，完全采用float32的数据精度会造成很大的存储空间浪费，这也为空间优化提供了一个方向。目前混合精度训练（AMP）即有三方的实现也有Pytorch的官方集成版本：

1. https://github.com/NVIDIA/apex
2. https://pytorch.org/docs/stable/amp.html

这里对官方提供的方法进行简单的介绍：(更详细的内容请参考官方文档）

• torch.cuda.amp.autocast(enabled=True)

autocast有上下文管理器（context manager）以及装饰器（decorator）两种使用方法

1. context manager

注意：文档中推荐在完成forward后退出上下文管理器，因为如果在BP的过程中仍然处于管理器中，那么在自动调节的机制下可能会出现类型不匹配错误，如下所述：

2. decorator

作为装饰器的使用方法十分简单，只需要在对应模型的FW的方法前进行装饰即可。

3. 手动恢复数据类型，避免数据不匹配错误

在我们对模型的精度需求很明确的情况下，为了避免退出管理器后出现精度匹配的错误，我们可以进行手动地恢复，下面再挂一段官方案例：

文档中的补充说明：在autocast管理器中我们是无需担心数据不匹配错误的，但并不一定不存在，也要注意。

4. autocast(enable=False)嵌套在autocast()中使用

这种方式是受推荐的，一些情况下我们可能想要特定的阶段运行在特定的数据类型精度下，那么只需要适时地关闭上下文管理器即可：

5. 自定义装饰器，指定数据类型
   使用以及分别针对forward以及backward pass，能够实现自定义的数据精度。使用方法就是装饰器。

• torch.cuda.amp.GradScaler

在float16精度下，当gradient的量级过小时可能会被忽略（归零）“下溢”，这对训练是十分不利的。GradScaler的思想是十分朴素的，即在进行backward之前，先对放大（scale）loss的值，直至避免或减少下溢发生。官方描述如下：

定义方法：

使用方法：

注意：如果enabled=False那么scalar不会起作用，而是直接调用optimizer.step()。

---

除了train阶段，在无需进行bcw的地方，不仅仅是从节省显存的角度考虑，也有助于我们更好地把控整个网络：如在validation阶段我们可以通过上下文管理器model.eval()，将各模型调整至evaluation阶段；或者直接通过with torch.no_grad()避免bwd行为。需要注意的是：model.eval()与with torch.no_grad()并不是等价的。具体而言：model.eval()会改变各网络层次的行为，如Dropout层，在train阶段与eval阶段的行为表现存在差异。可以参考：https://discuss.pytorch.org/t/model-eval-vs-with-torch-no-grad/19615

---

该命令的作用可以简述为：释放Pytorch的memory allocator所占用的“无用”的空间。特别需要注意的是“无用的”三个字，也就是说该命令是无法直接根据我们的需要释放指定的空间的，而由分配器统一管理。。

---

假设网络层次太深，那么可以将网络进行切割，先对前部分进行训练再对后部分进行训练。例如网络有100层，一次性训练时显存无法支持，那么就进行切割，将其分为网络1与网络2，分别训练。这里贴一段示例代码：注：代码来自 https://oldpan.me/archives/how-to-use-memory-pytorch

对于nn.Sequiential的情况则可以：

这里补充一段官方文档的注释：

从上图中可以看出：最后一个segment不会运行在‘torch.no_grad()’模式下，也就是说如果我们将segments设置为1，那么checkpoint_sequential将表现为无效。

---

即减小batchsize的另一种形式，假设当batchsize=64时显存溢出，那么我们可以设定batchsize为32，训练两次，累计误差再一次性backward。。示例：

• 上一篇: 旅行社
• 下一篇: 本市进一步推进非居住建筑改建宿舍型租赁住房有关工作