PyTorch手动梯度清零

简单介绍

Pytorch中,每个batch训练完后需要使用Variable.grad.zero_()进行梯度清零(其中Variable为变量名,性质为torch.tensor).在Pytorch中,之所以需要手动进行梯度清零,而不是选择自动清零,是因为这种方式可以让使用者自由选择梯度清零的时机,具有更高的灵活性.例如选择训练每N个batch后再进行梯度更新和清零,这相当于将原来的batch_size扩大为N×batch_size.因为原先是每个batch_size训练完后直接更新,而现在变为N个batch_size训练完才更新,相当于将N个batch_size合为了一组.这样可以让使用者使用较低的配置,跑较高的batch_size.

具体用法

第一种,直接梯度清零

上图中的用法为常规用法,每完成一个batch,更新一次梯度,完成一次训练.

第二种,batch叠加

上图中关键步骤有两步
第一步为:

loss = loss / accumulation_setps

其中accumulation_setps即是我在简单介绍中提到的N,用来控制梯度更新时机.由于损失函数均需要求平均,如果没有上图中的代码,相当于accumulation_setps个batch的损失值简单相加(类似于batch中m个数据的损失相加后并不求平均),这显然是不合适的.
第二步为:

if((i+1)%accumulation_setps) == 0optimizer.step()optimizer.zero_grad()

该步意义简单明了,即是每accumulation_setps个batch梯度清一次零,并完成一次训练.
两张图片均来源于一篇知乎问答,具体在哪找不到了.