pytorch累计梯度训练

背景

在深度学习的训练中，往往我们需要选择一个适中的batch值取获得最好的训练效果，batch不能太小或者太大。但是有时候我们的显卡内存有限，比如我们的显卡只能容纳一个batch,但是单batch的训练效果不稳定，该如何解决呢？

解决方案

我们可以累计N个batch的梯度之后再更新参数，而不是传统的每个batch都进行参数更新，这样可以等同于batch=N的训练；

代码

首先，需要明确以下两点：
1.optimizer.zero_grad()：会把所有的梯度置为0
2.loss.backward()：添加（add)新计算的梯度到旧的梯度上
累计梯度的训练中我们要注意学习率的大小；
代码一：
在iter_size循环之后累计梯度再除以iter_size，学习率就不要变化

optimizer.zero_grad()
loss = 0
minibatch_size = old_batch_size / iter_size
for i in range(iter_size):
    # output here as the size of minibatch
    loss += criterion(output, target_var)
loss = loss / iter_size
loss.backward()
optimizer.step()

代码二：
因为同样loss也除以了iter_size，所以学习率保持不变

optimizer.zero_grad()
loss_sum = 0
for  i in range(iter_size):
    loss = criterion(output, target_var) / iter_size
    loss.backward()
    loss_sum += loss
optimizer.step()

实验

从单batch->累计N batch之后，可以明显感受到网络更稳定更快地收敛了