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)  # 1loss = criterion(pred, label)  # 2optimizer.zero_grad()  # 3loss.backward()  # 4optimizer.step()  # 5

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

使用梯度累加的方法如下

for i,(image, label) in enumerate(train_loader):# 1. input outputpred = model(image)loss = criterion(pred, label)# 2.1 loss regularizationloss = loss / accumulation_steps  # 2.2 back propagationloss.backward()# 3. update parameters of netif (i+1) % accumulation_steps == 0:# optimizer the netoptimizer.step()        # update parameters of netoptimizer.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) / Kloss.backward()  # 积累梯度，但是多卡之间不进行同步# 第K个stepprediction = model(data)loss = loss_fn(prediction, label) / Kloss.backward()  # 进行多卡之间的梯度同步optimizer.step()

优雅写法

from contextlib import nullcontext
# 如果你的python版本小于3.7，请注释掉上面一行，使用下面这个：
# from contextlib import suppress as nullcontextif local_rank != -1:model = DDP(model)optimizer.zero_grad()
for i, (data, label) in enumerate(dataloader):# 只在DDP模式下，轮数不是K整数倍的时候使用no_syncmy_context = model.no_sync if local_rank != -1 and i % K != 0 else nullcontextwith my_context():prediction = model(data)loss = loss_fn(prediction, label) / Kloss.backward()  # 积累梯度，不应用梯度改变if i % K == 0:optimizer.step()optimizer.zero_grad()

梯度累加的影响

BN的影响