【CUDA out of memory】显存不足时的Trick

显存不足时的Trick

• 降低batch_size
  适当降低batch size， 则模型每层的输入输出就会成线性减少
• 数据类型
  32位的浮点数，切换到 16位的浮点数
• 梯度累积
  当执行 loss.backward() 时， 会为每个参数计算梯度，并将其存储在 paramter.grad 中， 注意到， paramter.grad 是一个张量， 其会累加每次计算得到的梯度。

在 Pytorch 中， 只有调用 optimizer.step()时才会进行梯度下降更新网络参数。

batch size 与占用显存息息相关，但有时候我们的batch size 又不能设置的太小，这咋办呢？答案就是梯度累加

传统训练：

for i,(feature,target) in enumerate(train_loader):
    outputs = model(feature)  # 前向传播
    loss = criterion(outputs,target)  # 计算损失
 
    optimizer.zero_grad()   # 清空梯度
    loss.backward()  # 计算梯度
    optimizer.step()  # 反向传播， 更新网络参数

加入梯度累加后：

for i,(features,target) in enumerate(train_loader):
    outputs = model(images)  # 前向传播
    loss = criterion(outputs,target)  # 计算损失
    loss = loss/accumulation_steps   # 可选，如果损失要在训练样本上取平均
 
    loss.backward()  # 计算梯度
    if((i+1)%accumulation_steps)==0:
        optimizer.step()        # 反向传播，更新网络参数
        optimizer.zero_grad()   # 清空梯度

本质上就是累加 accumulation_steps 个 batchsize/accumulationsteps 的梯度， 再根据累加的梯度来更新网络参数，以达到真实梯度类似batch_size 的效果

使用时，需要注意适当的扩大学习率

举例来说，假设 batch size = 4 ， accumulation steps = 8 ， 梯度积累首先在前向传播的时候以 batch_size=4 来计算梯度，但是不更新参数，将梯度积累下来，直到我们计算了 accumulation steps 个 batch， 我们再更新参数。其实本质上就等价于：

真正的 batch_size = batch_size * accumulation_steps