把显存用在刀刃上！17 种 pytorch 节约显存技巧

1. 显存都用在哪儿了？

一般在训练神经网络时，显存主要被网络模型和中间变量占用。

• 网络模型中的卷积层，全连接层和标准化层等的参数占用显存，而诸如激活层和池化层等本质上是不占用显存的。
• 中间变量包括特征图和优化器等，是消耗显存最多的部分。
• 其实 pytorch 本身也占用一些显存的，但占用不多，以下方法大致按照推荐的优先顺序。

2. 技巧 1：使用就地操作

就地操作 (inplace) 字面理解就是在原地对变量进行操作，对应到 pytorch 中就是在原内存上对变量进行操作而不申请新的内存空间，从而减少对内存的使用。具体来说就地操作包括三个方面的实现途径：

• 使用将 inplace 属性定义为 True 的激活函数，如 nn.ReLU(inplace=True)
• 使用 pytorch 带有就地操作的方法，一般是方法名后跟一个下划线 “_”，如 tensor.add_()，tensor.scatter_()，F.relu_()
• 使用就地操作的运算符，如 y += x，y *= x

3. 技巧 2：避免中间变量

在自定义网络结构的成员方法 forward 函数里，避免使用不必要的中间变量，尽量在之前已申请的内存里进行操作，比如下面的代码就使用太多中间变量，占用大量不必要的显存：

def forward(self, x):

	x0 = self.conv0(x)  # 输入层
	x1 = F.relu_(self.conv1(x0) + x0)
	x2 = F.relu_(self.conv2(x1) + x1)
	x3 = F.relu_(self.conv3(x2) + x2)
	x4 = F.relu_(self.conv4(x3) + x3)
	x5 = F.relu_(self.conv5(x4) + x4)
	x6 = self.conv(x5)  # 输出层

	return x6

为了减少显存占用，可以将上述 forward 函数修改如下：

def forward(self, x):

	x = self.conv0(x)  # 输入层
	x = F.relu_(self.conv1(x) + x)
	x = F.relu_(self.conv2(x) + x)
	x = F.relu_(self.conv3(x) + x)
	x = F.relu_(self.conv4(x) + x)
	x = F.relu_(self.conv5(x) + x)
	x = self.conv(x)  # 输出层

	return x

上述两段代码实现的功能是一样的，但对显存的占用却相去甚远，后者能节省前者占用显存的接近 90% 之多。

4. 技巧 3：优化网络模型

网络模型对显存的占用主要指的就是卷积层，全连接层和标准化层等的参数，具体优化途径包括但不限于：

• 减少卷积核数量 (=减少输出特征图通道数)
• 不使用全连接层
• 全局池化 nn.AdaptiveAvgPool2d() 代替全连接层 nn.Linear()
• 不使用标准化层
• 跳跃连接跨度不要太大太多 (避免产生大量中间变量)

5. 技巧 4：减小 BATCH_SIZE

• 在训练卷积神经网络时，epoch 代表的是数据整体进行训练的次数，batch 代表将一个 epoch 拆分为 batch_size 批来参与训练。
• 减小 batch_size 是一个减小显存占用的惯用技巧，在训练时显存不够一般优先减小 batch_size ，但 batch_size 不能无限变小，太大会导致网络不稳定，太小会导致网络不收敛。

6. 技巧 5：拆分 BATCH

拆分 batch 跟技巧 4 中减小 batch_size 本质是不一样的， 这种拆分 batch 的操作可以理解为将两次训练的损失相加再反向传播，但减小 batch_size 的操作是训练一次反向传播一次。拆分 batch 操作可以理解为三个步骤，假设原来 batch 的大小 batch_size=64：

• 将 batch 拆分为两个 batch_size=32 的小 batch
• 分别输入网络与目标值计算损失，将得到的损失相加
• 进行反向传播

7. 技巧 6：降低 PATCH_SIZE

• 在卷积神经网络训练中，patch_size 指的是输入神经网络的图像大小，即(H*W)。
• 网络输入 patch 的大小对于后续特征图的大小等影响非常大，训练时可能采用诸如 [64*64]，[128*128] 等大小的 patch，如果显存不足可以进一步缩小 patch 的大小，比如 [32*32]，[16*16]。
• 但这种方法存在问题，可能极大地影响网络的泛化能力，在裁剪的时候一定要注意在原图上随机裁剪，一般不建议。

8. 技巧 7：优化损失求和

一个 batch 训练结束会得到相应的一个损失值，如果要计算一个 epoch 的损失就需要累加之前产生的所有 batch 损失，但之前的 batch 损失在 GPU 中占用显存，直接累加得到的 epoch 损失也会在 GPU 中占用显存，可以通过如下方法进行优化：

epoch_loss += batch_loss.detach().item()  # epoch 损失

上边代码的效果就是首先解除 batch_loss 张量的 GPU 占用，将张量中的数据取出再进行累加。

9. 技巧 8：调整训练精度

• 降低训练精度
  pytorch 中训练神经网络时浮点数默认使用 32 位浮点型数据，在训练对于精度要求不是很高的网络时可以改为 16 位浮点型数据进行训练，但要注意同时将数据和网络模型都转为 16 位浮点型数据，否则会报错。降低浮点型数据的操作实现过程非常简单，但如果优化器选择 Adam 时可能会报错，选择 SGD 优化器则不会报错,具体操作步骤如下：

model.cuda().half()  # 网络模型设置半精度
# 网络输入和目标设置半精度
x, y = Variable(x).cuda().half(), Variable(y).cuda().half()

• 混合精度训练
  混合精度训练指的是用 GPU 训练网络时，相关数据在内存中用半精度做储存和乘法来加速计算，用全精度进行累加避免舍入误差，这种混合经度训练的方法可以令训练时间减少一半左右，也可以很大程度上减小显存占用。在 pytorch1.6 之前多使用 NVIDIA 提供的 apex 库进行训练，之后多使用 pytorch 自带的 amp 库，实例代码如下：

import torch
from torch.nn.functional import mse_loss
from torch.cuda.amp import autocast, GradScaler

EPOCH = 10  # 训练次数
LEARNING_RATE = 1e-3  # 学习率

x, y = torch.randn(3, 100).cuda(), torch.randn(3, 5).cuda()  # 定义网络输入输出
myNet = torch.nn.Linear(100, 5).cuda()  # 实例化网络，一个全连接层

optimizer = torch.optim.SGD(myNet.parameters(), lr=LEARNING_RATE)  # 定义优化器
scaler = GradScaler()  # 梯度缩放

for i in range(EPOCH):  # 训练

    with autocast():  # 设置混合精度运行
        y_pred = myNet(x)
        loss = mse_loss(y_pred, y)

    scaler.scale(loss).backward()  # 将张量乘以比例因子，反向传播
    scaler.step(optimizer)  # 将优化器的梯度张量除以比例因子。
    scaler.update()  # 更新比例因子

10. 技巧 9：分割训练过程

• 如果训练的网络非常深，比如 resnet101 就是一个很深的网络，直接训练深度神经网络对显存的要求非常高，一般一次无法直接训练整个网络。在这种情况下，可以将复杂网络分割为两个小网络，分别进行训练。
• checkpoint 是 pytorch 中一种用时间换空间的显存不足解决方案，这种方法本质上减少的是参与一次训练网络整体的参数量，如下是一个实例代码。

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

# 自定义函数
def conv(inplanes, outplanes, kernel_size, stride, padding):
    return nn.Sequential(nn.Conv2d(inplanes, outplanes, kernel_size, stride, padding),
                         nn.BatchNorm2d(outplanes),
                         nn.ReLU()
                         )


class Net(nn.Module):  # 自定义网络结构，分为三个子网络
    def __init__(self):
        super().__init__()

        self.conv0 = conv(3, 32, 3, 1, 1)
        self.conv1 = conv(32, 32, 3, 1, 1)
        self.conv2 = conv(32, 64, 3, 1, 1)
        self.conv3 = conv(64, 64, 3, 1, 1)
        self.conv4 = nn.Linear(64, 10)  # 全连接层

    def segment0(self, x):  # 子网络1
        x = self.conv0(x)
        return x

    def segment1(self, x):  # 子网络2
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        return x

    def segment2(self, x):  # 子网络3
        x = self.conv4(x)
        return x

    def forward(self, x):

        x = checkpoint(self.segment0, x)  # 使用 checkpoint
        x = checkpoint(self.segment1, x)
        x = checkpoint(self.segment2, x)

        return x

• 使用 checkpoint 进行网络训练要求输入属性 requires_grad=True ，在给出的代码中将一个网络结构拆分为 3 个子网络进行训练，对于没有 nn.Sequential() 构建神经网络的情况无非就是自定义的子网络里多几项，或者像例子中一样单独构建网络块。
• 对于由 nn.Sequential() 包含的大网络块 (小网络块时没必要)，可以使用 checkpoint_sequential 包来简化实现，具体实现过程如下：

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint_sequential


class Net(nn.Module):  # 自定义网络结构，分为三个子网络
    def __init__(self):
        super().__init__()
		linear = [nn.Linear(10, 10) for _ in range(100)]
        self.conv = nn.Sequential(*linear)  # 网络主体，100 个全连接层

    def forward(self, x):

        num_segments = 2  # 拆分为两段
        x = checkpoint_sequential(self.conv, num_segments, x)

        return x

11. 技巧10：清理内存垃圾

• python 中定义的变量一般在使用结束时不会立即释放资源，在训练循环开始时可以利用如下代码来回收内存垃圾。

import gc 
gc.collect()  # 清理内存

12. 技巧11：使用梯度累积

• 由于显存大小的限制，训练大型网络模型时无法使用较大的 batch_size ，而一般较大的 batch_size 能令网络模型更快收敛。
• 梯度累积就是将多个 batch 计算得到的损失平均后累积再进行反向传播，类似于技巧 5 中拆分 batch 的思想(但技巧 5 是将大 batch 拆小，训练的依旧是大 batch，而梯度累积训练的是小 batch)。
• 可以采用梯度累积的思想来模拟较大 batch_size 可以达到的效果，具体实现代码如下：

output = myNet(input_)  # 输入送入网络
loss = mse_loss(target, output)  # 计算损失
loss = loss / 4  # 累积 4 次梯度
loss.backward()  # 反向传播
if step % 4 == 0:  # 如果执行了 4 步
    optimizer.step()  # 更新网络参数
    optimizer.zero_grad()  # 优化器梯度清零

13. 技巧12：清除不必要梯度

在运行测试程序时不涉及到与梯度有关的操作，因此可以清楚不必要的梯度以节约显存，具体包括但不限于如下操作：

• 用代码 model.eval() 将模型置于测试状态，不启用标准化和随机舍弃神经元等操作。
• 测试代码放入上下文管理器 with torch.no_grad(): 中，不进行图构建等操作。
• 在训练或测试每次循环开始时加梯度清零操作

myNet.zero_grad()  # 模型参数梯度清零
optimizer.zero_grad()  # 优化器参数梯度清零

14. 技巧13：周期清理显存

• 同理也可以在训练每次循环开始时利用 pytorch 自带清理显存的代码来释放不用的显存资源。

 torch.cuda.empty_cache()  # 释放显存

执行这条语句释放的显存资源在用 Nvidia-smi 命令查看时体现不出，但确实是已经释放。其实 pytorch 原则上是如果变量不再被引用会自动释放，所以这条语句可能没啥用，但个人觉得多少有点用。

15. 技巧14：多使用下采样

下采样从实现上来看类似池化，但不限于池化，其实也可以用步长大于 1 来代替池化等操作来进行下采样。从结果上来看就是通过下采样得到的特征图会缩小，特征图缩小自然参数量减少，进而节约显存，可以用如下两种方式实现：

nn.Conv2d(32, 32, 3, 2, 1)  # 步长大于 1 下采样

nn.Conv2d(32, 32, 3, 1, 1)  # 卷积核接池化下采样
nn.MaxPool2d(2, 2)

16. 技巧15：删除无用变量

del 功能是彻底删除一个变量，要再使用必须重新创建，注意 del 删除的是一个变量而不是从内存中删除一个数据，这个数据有可能也被别的变量在引用，实现方法很简单，比如：

 def forward(self, x):
 
	input_ = x
	x = F.relu_(self.conv1(x) + input_)
	x = F.relu_(self.conv2(x) + input_)
	x = F.relu_(self.conv3(x) + input_)
	
	del input_  # 删除变量 input_

	x = self.conv4(x)  # 输出层
	return x

17. 技巧16：改变优化器

进行网络训练时比较常用的优化器是 SGD 和 Adam，抛开训练最后的效果来谈，SGD 对于显存的占用相比 Adam 而言是比较小的，实在没有办法时可以尝试改变参数优化算法，两种优化算法的调用是相似的：

import torch.optim as optim
from torchvision.models import resnet18

LEARNING_RATE = 1e-3  # 学习率
myNet = resnet18().cuda()  # 实例化网络

optimizer_adam = optim.Adam(myNet.parameters(), lr=LEAENING_RATE)  #  adam 网络参数优化算法
optimizer_sgd = optim.SGD(myNet.parameters(), lr=LEAENING_RATE)  # sgd 网络参数优化算法

18. 终极技巧

购买显存够大的显卡，一块不行那就 多来几块。