pytorch梯度累积

pytorch梯度累积

梯度累积技术是为了解决显存不足的问题

for _,(x,y) in enumerate(data_loader):
	optimizer.zero_grad() # 梯度清零
  y_pred=model(x) # 预测
  loss=loss_fn(y,y_pred) # 计算损失
  loss.backward() # 反向传播
  optimizer.step()

pytorch每次前向传播后得到反向传播的计算图，其中变量的梯度不会清空，如果没有optimizer.zero_grad(),将会对变量的梯度进行累积。

accumulation_step=4
for i,(x,y) in enumerate(data_loader):
  y_pred=model(x) # 预测
  loss=loss_fn(y,y_pred) # 计算损失
  loss.backward() # 反向传播
  if (i+1) % accumulation_steps == 0:
  	optimizer.step()
  	optimizer.zero_grad() # 梯度清零

设置了accumulation_steps = 4，只有当小batch运算四次的时候才会进行清空梯度，更新参数，相当于扩大了4倍batch_size

讲一下pytorch梯度累积对bn层影响的看法

batchnormalization在训练阶段会对已经跑过的batch进行均值和方差的计算，计算公式如下

$running\_mean=(1-momentum)\ast running\_mean+momentum\ast \mu$

$running\_var=(1-momentum)\ast running\_var+momentum\ast \sigma$

其中momentun默认值为0.1，当进行梯度累积的时候，我们每次forward都会进行$mean和var$的计算，导致记忆的历史mean和var的视野会缩短，更久远的统计信息占比的权重会更小，所以相比于大batch,梯度累积方式的bn层计算的$mean$和$var$会不那么精确，可以调低momentum的值来增强记忆能力，以逼近真实的扩大batch_size的效果

参考文献：

• pytorch骚操作之梯度累加，变相增大batch size_如何把梯度叠加后取平均-CSDN博客

• Pytorch-BN层详细解读_pytorch bn层-CSDN博客