---

前言

如今的深度学习发展的如火如荼，相信各行各业的大家都或多或少接触过深度学习的知识。相信很多人在跑模型时都见过以下语句：
RuntimeError: CUDA out of memory.
显存不足是很多人感到头疼的问题，毕竟能拥有大量显存的实验室还是少数，而现在的模型已经越跑越大，模型参数量和数据集也越来越大。在别人的Batchsize在32甚至64以上时，我们还在为强行塞下1个Batch而努力。

以下给大家提供一些节省PyTorch显存占用的小技巧，虽然提升不大，但或许能帮你达到可以勉强运行的及格线。

一、大幅减少显存占用方法

想大幅减少显存占用，必定要从最占用显存的方面进行缩减，即 模型 和 数据。

1. 模型

在模型上主要是将Backbone改用轻量化网络或者减少网络层数等方法，可以很大程度上减少模型参数量，从而减少显存占用。

2. 数据

数据方面上要减少显存占用主要是使用小 BatchSize 或者将输入数据 Resize 到较小的尺寸。

二、小幅减少显存占用方法

有时候我们可能不想更改模型，而又恰好差一点点显存或者想尽量多塞几个BatchSize，有一些小技巧可以挤出一点点显存。

1. 使用inplace

PyTorch中的一些函数，例如 ReLU、LeakyReLU 等，均有 inplace 参数，可以对传入Tensor进行就地修改，减少多余显存的占用。

2. 加载、存储等能用CPU就绝不用GPU

GPU存储空间宝贵，我们可以选择使用CPU做一些可行的分担，虽然数据传输会浪费一些时间，但是以时间换空间，可以视情况而定，在模型加载中，如 torch.load_state_dict 时，先加载再使用 model.cuda()，尤其是在 resume 断点续训时，可能会报显存不足的错误。数据加载也是，在送入模型前在送入GPU。其余中间的数据处理也可以依循这个原则。

3. 低精度计算

可以使用 float16 半精度混合计算，也可以有效减少显存占用，但是要注意一些溢出情况，如 mean 和 sum等。

4. torch.no_grad

对于 eval 等不需要 bp 及 backward 的时候，可已使用with torch.no_grad，这个和model.eval()有一些差异，可以减少一部分显存占用。

5. 及时清理不用的变量

对于一些使用完成后的变量，及时del掉，例如 backward 完的 Loss，缓存torch.cuda.empty_cache()等。

6. 分段计算

骚操作，我们可以将模型或者数据分段计算。

1. 模型分段，利用checkpoint将模型分段计算

   # 首先设置输入的input=>requires_grad=True
   # 如果不设置可能会导致得到的gradient为0
   input = torch.rand(1, 10, requires_grad=True)
   layers = [nn.Linear(10, 10) for _ in range(1000)]
   
   # 定义要计算的层函数，可以看到我们定义了两个
   # 一个计算前500个层，另一个计算后500个层
   def run_first_half(*args):
       x = args[0]
       for layer in layers[:500]:
           x = layer(x)
       return x
   
   def run_second_half(*args):
       x = args[0]
       for layer in layers[500:-1]:
           x = layer(x)
       return x
   
   # 引入checkpoint
   from torch.utils.checkpoint import checkpoint
   
   x = checkpoint(run_first_half, input)
   x = checkpoint(run_second_half, x)
   # 最后一层单独执行
   x = layers[-1](x)
   x.sum.backward()
   

2. 数据分段，例如原来需要64个batch的数据forward一次后backward一次，现在改为32个batch的数据forward两次后backward一次。

总结

以上是我总结的一些PyTorch节省显存的一些小技巧，希望可以帮助到大家，如果有其它好方法，也欢迎和我讨论。