Pytorch并行计算(三): 梯度累加

梯度累加

梯度累加（Gradient Accmulation）是一种增大训练时batch size的技巧。当batch size在一张卡放不下时，可以将很大的batch size分解为一个个小的mini batch，分别计算每一个mini batch的梯度，然后将其累加起来优化

正常的pytorch训练流程如下（来自知乎）

for i, (image, label) in enumerate(train_loader):
    pred = model(image)  # 1
    loss = criterion(pred, label)  # 2
    optimizer.zero_grad()  # 3
    loss.backward()  # 4
    optimizer.step()  # 5

1. 神经网络forward过程
2. 获取loss，通过pred和label计算你损失函数
3. 清空网络中参数的梯度
4. 反向传播，计算当前梯度
5. 根据梯度更新网络参数

使用梯度累加的方法如下

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

1. 神经网络forward过程，同时计算损失函数
2. 反向传播计算当前梯度（在backward时，计算的loss要除batch的大小得到均值）
3. 不断重复1、2步骤，重复获取梯度
4. 梯度累加到一定次数后，先optimizer.step()更新网络参数，随后zero_grad()清除梯度，为下一次梯度累加做准备

DDP中的梯度累加

问题：在DDP中所有卡的梯度all_reduce阶段发生在loss.bachward()阶段，也就是说执行loss.backward()之后，所有卡的梯度会进行一次汇总，但是如果我们如果使用梯度累加策略，假设梯度累加K=2，就需要all_reduce汇总两次，会带来额外的计算错误和时间开销

解决方案：知乎写的很好，这里参考其解决方案，只需要在前K-1次取消梯度同步即可，DDP提供了一个暂时取消梯度同步的context函数no_sync()，在这个函数下，DDP不会进行梯度同步

model = DDP(model)

for 每次梯度累加循环
    optimizer.zero_grad()
    # 前accumulation_step-1个step，不进行梯度同步，每张卡分别累积梯度。
    for _ in range(K-1)::
        with model.no_sync():
            prediction = model(data)
            loss = loss_fn(prediction, label) / K
            loss.backward()  # 积累梯度，但是多卡之间不进行同步
    # 第K个step
    prediction = model(data)
    loss = loss_fn(prediction, label) / K
    loss.backward()  # 进行多卡之间的梯度同步
    optimizer.step()

优雅写法

from contextlib import nullcontext
# 如果你的python版本小于3.7，请注释掉上面一行，使用下面这个：
# from contextlib import suppress as nullcontext

if local_rank != -1:
    model = DDP(model)

optimizer.zero_grad()
for i, (data, label) in enumerate(dataloader):
    # 只在DDP模式下，轮数不是K整数倍的时候使用no_sync
    my_context = model.no_sync if local_rank != -1 and i % K != 0 else nullcontext
    with my_context():
        prediction = model(data)
        loss = loss_fn(prediction, label) / K
        loss.backward()  # 积累梯度，不应用梯度改变
    if i % K == 0:
        optimizer.step()
        optimizer.zero_grad()

梯度累加的影响

BN的影响