pytorch减少显存方式

标签：显存技巧 nn self batch pytorch 训练减少

引导

1. 显存都用在哪儿了？
2. 技巧 1：使用就地操作
3. 技巧 2：避免中间变量
4. 技巧 3：优化网络模型
5. 技巧 4：减小 BATCH_SIZE
6. 技巧 5：拆分 BATCH
7. 技巧 6：降低 PATCH_SIZE
8. 技巧 7：优化损失求和
9. 技巧 8：调整训练精度
10. 技巧 9：分割训练过程
11. 技巧10：清理内存垃圾
12. 技巧11：使用梯度累积
13. 技巧12：清除不必要梯度
14. 技巧13：周期清理显存
15. 技巧14：多使用下采样
16. 技巧15：删除无用变量
17. 技巧16：改变优化器
18. 终极技巧

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. 终极技巧

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

标签：显存,技巧,nn,self,batch,pytorch,训练,减少
From： https://www.cnblogs.com/chentiao/p/17901639.html

引导