PyTorch 梯度累积小技巧

受限于 GPU 内存的大小，在训练神经网络的时候，我们经常会在不改变其他的训练参数的情况下（比如 learning rate），用多次小的 mini-batch 来模拟一个较大的 mini-batch，即：

$$
\begin{align}
\mbox{global_batch_size}&=\mbox{batch_size}*\mbox{iter_size}\
\end{align}
$$

先给出两种不同的做法。下面是 naive 的一种写法：

// "blah-blah"
optimizer.zero_grad()
loss = 0
# minibatch_size = global_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()

第二种做法：

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

对比一下两种做法：

方法一将 loss 叠加多次，这样的做法会导致计算图的中间变量会一直保存，占用内存。
方法二是进行多次误差反向传播，并对梯度进行累加，每次反传 (backward) 都会清除计算图中的中间变量，这样会节省内存。另外，还值得注意的是，为了经过 iter_size 次迭代后的累积梯度跟使用一个大的batch_size*iter_size 一样，需要将每次的 loss 除以 iter_size。
附带讲，如果想在反向传播计算后还保存计算图中的中间变量，可以指定 retain_graph 选项，即 loss.backward(retain_graph=True)。

Don’t forget to divide the loss by iter_size to normalize the gradient. The second version will give you same result as if you are having larger mini-batch size. ^[1]

注意的是，累积梯度的做法，不适用于模拟大 batchsize 的 Batch Normalization。

1.PyTorch Forums: How to implement accumulated gradient？ ↩