当 GPU 的内存不够时，如何使用大批量（large batch size）样本来训练神经网络呢？

解决方法 -- 梯度累加（accumulating gradients）

 

一般在 PyTorch 中，我们是这样来更新梯度的：

在计算梯度的 loss.backward() 操作中，每个参数的梯度被计算出来后，都被存储在各个参数对应的一个张量里：parameter.grad。然后优化器就会根据这个来更新每个参数的值，就是 optimizer.step()。

 

而梯度累加（accumulating gradients）的基本思想就是， 在优化器更新参数前，也就是执行 optimizer.step() 前，我们进行多次梯度计算，保存在 parameter.grad 中，然后累加梯度再更新。这个在 PyTorch 中特别容易实现，因为 PyTorch 中，梯度值本身会保留，除非我们调用 model.zero_grad() or optimizer.zero_grad()。

下面是一个梯度累加的例子，其中 accumulation_steps 就是要累加梯度的循环数：

 

如果样本特别大，别说 batch training，要是 GPU 的内存连一个样本都不下怎么办呢？

答案是使用梯度检查点（gradient-checkpoingting），用计算量来换内存。基本思想就是，在反向传播的过程中，把梯度切分成几部分，分别对网络上的部分参数进行更新（见下图）。但这种方法的速度很慢，因为要增加额外的计算量。