梯度累加(Gradient Accumulation)

1,270 阅读2分钟

我们在训练神经网络的时候,超参数batch size的大小会对最终的模型效果产生很大的影响。一定条件下,batch size设置的越大,模型就会越稳定。batch size的值通常设置在 8-32 之间,但是当我们做一些计算量需求大的任务(例如语义分割、GAN等)或者输入图片尺寸太大的时候,我们的batch size往往只能设置为2或者4,否则就会出现 “CUDA OUT OF MEMORY” 的不可抗力报错。

贫穷是促进人类进步的阶梯,如何在有限的计算资源的条件下,训练时采用更大的batch size呢?这就是梯度累加(Gradient Accumulation)技术了。

我们以Pytorch为例,一个神经网络的训练过程通常如下:

for i, (inputs, labels) in enumerate(trainloader):
    optimizer.zero_grad()                   # 梯度清零
    outputs = net(inputs)                   # 正向传播
    loss = criterion(outputs, labels)       # 计算损失
    loss.backward()                         # 反向传播,计算梯度
    optimizer.step()                        # 更新参数
if (i+1) % evaluation_steps == 0:
    evaluate_model()

从代码中可以很清楚地看到神经网络是如何做到训练的:
1.将前一个batch计算之后的网络梯度清零
2.正向传播,将数据传入网络,得到预测结果
3.根据预测结果与label,计算损失值
4.利用损失进行反向传播,计算参数梯度
5.利用计算的参数梯度更新网络参数

下面来看梯度累加是如何做的:

model.zero_grad()                                   # 梯度清零
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # 正向传播
    loss = loss_function(predictions, labels)       # 计算损失
    loss = loss / accumulation_steps                # Normalize our loss (if averaged)
    loss.backward()                                 # 反向传播,计算梯度
    if (i+1) % accumulation_steps == 0:             # 经过几次accumulation_steps
        optimizer.step()                            # 更新参数
        model.zero_grad()                           # 梯度清零
        if (i+1) % evaluation_steps == 0:           # Evaluate the model when we...
            evaluate_model()                        # ...have no gradients accumulated

1.正向传播,将数据传入网络,得到预测结果
2.根据预测结果与label,计算损失值
3.利用损失进行反向传播,计算参数梯度
4.重复1-3,不清空梯度,而是将梯度累加
5.梯度累加达到固定次数之后,更新参数,然后将梯度清零

总结来讲,梯度累加就是每计算一个batch的梯度,不进行清零,而是做梯度的累加,当累加到一定的次数(accumulation_steps)之后,再更新网络参数,然后将梯度清零。

通过这种参数延迟更新的手段,可以实现与采用大batch size相近的效果。在平时的实验过程中,我一般会采用梯度累加技术,大多数情况下,采用梯度累加训练的模型效果,要比采用小batch size训练的模型效果要好很多。

image.png