梯度累计（Gradient Accumulation）是一种在显存受限情况下模拟大批次（large batch size）训练的技术。它的核心思想是：用多次小 batch 的前向/反向计算，累积梯度，对梯度除以 batch size得到平均梯度，然后一次性更新模型参数，从而在不增加显存占用的前提下，获得大 batch 训练的优化效果。
通常，深度学习训练中：前向传播 + 反向传播 会为一个 batch 计算梯度；优化器立即用该梯度更新模型参数。但当 batch size 很大时，中间激活和梯度会占用大量显存，可能超出 GPU 显存。当batch size很小，比如为1时不会超过GPU显存，它会每训练一个样本，然后计算一次梯度，由于样本之间的差异很大，导致更新的梯度忽大忽小不可控，导致训练过程不稳定、收敛缓慢，甚至无法收敛。
例如，下面是一个前向传播的计算图：
其中，a,b,c都是参数，他们会在反向传播的过程中会进行更新。更新的过程如下：
通过链式法则，可以得到损失函数对参数b的梯度，然后b-lr*该梯度作为新的b，再继续进行前向传播、反向传播更新，a和c同理。在这个过程中，模型的参数例如a,b,c等，称为叶子节点张量，其梯度会被累积到.grad属性中，并长期驻留现存，直到下一次optimizer.zero_grad()清空，对于中间非叶子张量，例如v，其值等于v=b*c，其梯度默认不会保留，计算完后立即释放。
首先准备训练集和标签：
x,y=sklearn.datasets.load_digits(return_X_y=True)
x=torch.tensor(x/16).float().cuda() # FP32
y=torch.tensor(y).long().cuda()
print(x.shape,x.dtype)
print(y.shape,y.dtype)定义一个模型：
class MLP(torch.nn.Module):
def __init__(self,input_size,hidden_size,output_size):
super(MLP,self).__init__()
self.fc1=torch.nn.Linear(input_size, hidden_size)
self.fc2=torch.nn.Linear(hidden_size, output_size)
def forward(self,x):
out=self.fc1(x)
out=torch.relu(out)
out=self.fc2(out)
return out接下来实现梯度累计：
model=MLP(input_size=64,hidden_size=256,output_size=10).cuda()
loss_fn=torch.nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)
iter=0
accum_steps=0
while True:
out=model(x)
loss=loss_fn(out,y)
loss=loss/4
loss.backward()
accum_steps+=1
if accum_steps==4:
optimizer.step()
optimizer.zero_grad()
accum_steps=0
iter+=1
if iter%25000==0:
print(f'iter={iter} loss={loss.item()} cuda_mem={torch.cuda.memory_allocated()}Bytes')
if loss.item()<=1e-3:
breakaccum_steps用于梯度累计计数，其值为4时，更新一次参数。需要注意loss=loss/4，这是因为前面累积了4个batch的梯度，在更新的时候，应该要取batch size的均值进行更新。使用loss=loss/4并不影响计算图，可以理解为在loss后面增加new_loss=loss*(1/4)，当求梯度d_new_loss/d_loss时，值就是1/4，从而实现对loss乘以1/4。