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

Posted 桥上风景窗前人

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了把显存用在刀刃上!17 种 pytorch 节约显存技巧相关的知识,希望对你有一定的参考价值。

1. 显存都用在哪儿了?

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

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

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

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

  • 使用将 inplace 属性定义为 True 的激活函数,如 nn.ReLU(inplace=True)
  • 使用 pytorch 带有就地操作的方法,一般是方法名后跟一个下划线 “_”,如 tensor.add_()tensor.scatter_()F.relu_()
  • 使用就地操作的运算符,如 y += xy *= 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. 终极技巧

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

以上是关于把显存用在刀刃上!17 种 pytorch 节约显存技巧的主要内容,如果未能解决你的问题,请参考以下文章

显存不足?PyTorch 显存使用分析与优化

Azure手把手系列3:把IT的钱花在刀刃上

[Pytorch]深度模型的显存计算以及优化

PyTorch 半精度训练踩坑

PyTorch显存机制分析

pytorch显存越来越多的一个自己没注意的原因