我们在训练神经网络的时候,超参数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.利用计算的参数梯度更新网络参数

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

    for i, (inputs, labels) in enumerate(trainloader):

        outputs = net(inputs)                   # 正向传播
        loss = criterion(outputs, labels)       # 计算损失函数
        loss = loss / accumulation_steps        # 损失标准化
        loss.backward()                         # 反向传播,计算梯度

        if (i+1) % accumulation_steps == 0:
            optimizer.step()                    # 更新参数
            optimizer.zero_grad()               # 梯度清零

            if (i+1) % evaluation_steps == 0:
                evaluate_model()


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

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

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

参考资料

1.Loss Normalization
2.Model.zero_grad() or optimizer.zero_grad()?
3.A trick to use bigger batches for training: gradient accumulation
4.PyTorch中在反向传播前为什么要手动将梯度清零?
5.Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups



pytorch

本博客所有文章除特别声明外,均采用 CC BY-SA 3.0协议 。转载请注明出处!