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

本文所述的内存优化方法不包含例如“减少batch size”等直接影响训练流的方法

主要参考资料：部分其他资料将在文中直接给出

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

---

1.尽量使用in_place实现

使用in_place操作使得Pytorch的allocator不会记录该部分的原tensor，从而减少显存的消耗。也正是因为如此，如果在网络反向计算梯度的过程中需要用到原tensor，那么就不能够使用in_place操作，否则会使得反向传播报错。
ReLU：
Pytorch的nn.ReLU支持in_place操作，一般原则是，只要ReLU的前一层“不需要”进行梯度计算，那么ReLU就可以使用in_place实现：例如，当ReLU层前是一个BN层时（Pytorch的BN会在反向传播过程中重新计算对应的需求项，因此无需记录梯度信息）。

---

2.Pytorch的 ‘checkpoint’

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

model = nn.Sequential(nn.Conv2d(3,6,1,1), nn.ReLU()).to('cuda')
dummy_inp = torch.randn(1,3,128,128).to('cuda')
dummy_inp = autograd.Variable(dummy_input, requires_grad=True)
out = checkpoint_sequential(model, 2, dummy_input)
# 注意，如果这里的dummy_input没有requires_grad=True，会报错，因为checkpoint一定是在需要求梯度的情况下使用的，否则本来就不需要保存原值。

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

---

3.删除loss等不再被需要的值，释放对应的空间

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

for i, (x, y) in enumerate(train_loader):
    x = Variable(x)
    y = Variable(y)
    # compute model and update
    del x, y, output 

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

---

4.混合精度训练

一般而言，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

# 定义模型以及优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # 上下文管理器启动
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # 在backward阶段退出上下文管理器
    loss.backward()
    optimizer.step()

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

2. decorator

class AutocastModel(nn.Module):
    ...
    @autocast()
    def forward(self, input):
        ...

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

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

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

a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    # 假设torch.mm需求的是torch16精度的数据输入，输出同样为F16
    # 由于管理器的存在，我们不需要手动进行数据类型转换
    e_float16 = torch.mm(a_float32, b_float32)
    f_float16 = torch.mm(d_float32, e_float16)

# 这里已经退出了上下文管理器，所以需要手动进行类型转换
# 注意，虽然torch.mm运行在float16下，但是输入32的数据精度是没有问题的
g_float32 = torch.mm(d_float32, f_float16.float())

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

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

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

a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    e_float16 = torch.mm(a_float32, b_float32)

    with autocast(enabled=False):
        # 保证运行在float32的精度下
        f_float32 = torch.mm(c_float32, e_float16.float())
        
    g_float16 = torch.mm(d_float32, f_float32)

5. 自定义装饰器，指定数据类型
   使用torch.cuda.amp.custom_fwd(fwd=None, **kwargs)以及torch.cuda.amp.custom_bwd(fwd=None, **kwargs)分别针对forward以及backward pass，能够实现自定义的数据精度。使用方法就是装饰器。需要注意的是：这两个装饰器只有当位于autocast()上下文管理器中才能够发挥作用，否则不产生任何行为。

• torch.cuda.amp.GradScaler

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

定义方法：

GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

使用方法：

# 训练阶段开始前创建GradScalar
 |      scaler = GradScaler()
 |
 |      for epoch in epochs:
 |          for input, target in data:
 |              optimizer.zero_grad()
 |              output = model(input)
 |              loss = loss_fn(output, target)
 |
 |              # 先对loss进行放大，在backward
 |              scaler.scale(loss).backward()
 |
 				# 用sclar包装optimizer
 				# scaler首先unscale各parameter的gradient值
 				# 注意：如果缩小后的梯度值中没有产生上溢（inf/NaNs），那么接着调用optimizer.step()，否则跳过optimizer.step()
 |              scaler.step(optimizer)
 |
 |              # 依据定义的参数，更新scale值
 				# 这里的更新依据是在放缩后是否出现了上溢
 				# 如果有上溢那么backoff_factor起作用
 				# 如果没有上溢那么在growth_interval个轮次之后，growth_factor起作用。
 |              scaler.update()

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

---

5.model.eval()以及with torch.no_grad()

除了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

---

6.torch.cuda.empty_cache()

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

---

7.分割训练过程（基于checkpoint）

假设网络层次太深，那么可以将网络进行切割，先对前部分进行训练再对后部分进行训练。例如网络有100层，一次性训练时显存无法支持，那么就进行切割，将其分为网络1与网络2，分别训练。这里贴一段示例代码：注：代码来自 https://oldpan.me/archives/how-to-use-memory-pytorch
这里需要用到我们前面介绍的checkpoint特性, 机制我们在前面也已经进行介绍了

input = torch.rand(1, 10)
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)
# 方式一：
# output = model(input)

# 进行如下更改
# 设置输入的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()

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

input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]
model = nn.Sequential(*layers)

from torch.utils.checkpoint import checkpoint_sequential

# 分成两个部分
num_segments = 2
x = checkpoint_sequential(model, num_segments, input)
x.sum().backward() 

注意：该方法会影响与batchsize相关的层，例如batchNorm层
这里补充一段官方文档的注释：

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

---

8.分批训练，一次backward（loss accumulation）

即减小batchsize的另一种形式，假设当batchsize=64时显存溢出，那么我们可以设定batchsize为32，训练两次，累计误差再一次性backward。需要注意：该方法会影响与batchsize大小相关的层次，例如BatchNorm层。示例：

# loss accumulation
    accumulation_step = 4
    for i in range(K):
        train_dataset = tracking_dataset(kitti_object, root_dir=train_root, ki=i, K=K, typ='train')
        val_dataset = tracking_dataset(kitti_object, root_dir=train_root, ki=i, K=K, typ='val')
        train_loader = DataLoader(train_dataset, batch_size=2, drop_last=True)
        val_loader = DataLoader(val_dataset, batch_size=2, drop_last=True)
        for epoch in range(50):
            for idx, batch in enumerate(train_loader):
                print("batch{}.....".format(idx))
                for key in batch:
                    batch[key] = batch[key].to(device)
				inp = batch['inp']
                # output
                with autocast(enabled=True):
                    output = model(inp)
                    # loss
                    total_loss, losses = lossModel(output, batch)
                    loss_mean = total_loss.mean()
                loss_mean /= accumulation_step
                # accumulate
                loss_mean.backward()
                if (idx+1) % accumulation_step == 0:
                	# 一定次数后，计算梯度
                    optimizer.zero_grad()
                    optimizer.step()
                    print('success !!')