为什么Pytorch在反向传播前需要手动将梯度清零？

1 问题描述

日常开发工作中我们经常使用Pytorch深度学习框架进行模型的开发，利用Pytorch框架编写模型训练代码的一般流程如下：

for i, (image, label) in enumerate(train_loader):
    # 1. input output
    pred = model(image)
    loss = criterion(pred, label)
    # 2. backward
    optimizer.zero_grad()   # reset gradient
    loss.backward()
    optimizer.step()

上述代码解释如下：

• 获取输入数据：从dataloader中取出每个batch的输入数据，图像和标签
• 计算loss：输入图像送入网络，执行前向预测得到预测值，计算损失函数
• optimizer.zero_grad() 手动清空过往梯度
• loss.backward() 反向传播，根据loss值计算当前梯度
• optimizer.step() 根据梯度更新网络参数

简单来说，就是进来一个batch的数据，计算一次梯度，更新一次网络参数。
这里我们在网络回传之前，自己手动将梯度清零，这样的设计有什么好处呢？

2 梯度累加

一般来说，Pytorch针对梯度这样设计的目的在于在反向传播时，可以支持多种传播方式，玩出更多的花样，比如可通过以下方式实现梯度累加：

for i,(image, label) in enumerate(train_loader):
    # 1. input output
    pred = model(image)
    loss = criterion(pred, label)

    # 2.1 loss regularization
    loss = loss / accumulation_steps  

    # 2.2 back propagation
    loss.backward()

    # 3. update parameters of net
    if (i+1) % accumulation_steps == 0:
        # optimizer the net
        optimizer.step()        # update parameters of net
        optimizer.zero_grad()   # reset gradient

上述代码解释如下：

• 获取输入数据：从dataloader中取出每个batch的输入数据，图像和标签
• 计算loss：输入图像送入网络，执行前向预测得到预测值，计算损失函数
• loss.backward() 反向传播，计算当前梯度
• 多次循环步骤 1-2，不清空梯度，使梯度累加在已有梯度上
• 梯度累加了一定次数后，先optimizer.step() 根据累计的梯度更新网络参数，然后optimizer.zero_grad() 清空过往梯度，为下一波梯度累加做准备

简单来说：梯度累加就是，每次获取1个batch的数据，计算1次梯度，梯度不清空，不断累加，累加一定次数后，根据累加的梯度更新网络参数，然后清空梯度，进行下一次循环。

3 总结

一定条件下，batchsize 越大训练效果越好，梯度累加则实现了 batchsize 的变相扩大，如果accumulation_steps 为 8，则batchsize 相当于扩大了8倍。这样对于计算资源不是很充足的实验环境下解决显存受限的一个不错的trick，使用时需要注意，学习率也应该适当的放大。

4 附录

本文参考链接如下：
链接一

关注公众号《AI算法之道》，获取更多AI算法资讯。