转载自（侵删）：
https://www.cnblogs.com/chentiao/p/17901639.html
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 x11. 技巧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 x17. 技巧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 网络参数优化算法