首先要明确在Pytorch当中,计算得到的梯度是默认累加的,而不是下次计算梯度就自动清零上一次的梯度值。
这样做的好处有以下几点:
1、减小multitask的内存消耗
在PyTorch中,multi-task任务一个标准的train from scratch流程为:
for idx, data in enumerate(train_loader):
xs, ys = data
pred1 = model1(xs)
pred2 = model2(xs)
loss1 = loss_fn1(pred1, ys)
loss2 = loss_fn2(pred2, ys)
******
loss = loss1 + loss2
optmizer.zero_grad()
loss.backward()
++++++
optmizer.step()
从PyTorch的设计原理上来说,在每次进行前向计算得到pred时,会产生一个用于梯度回传的计算图,这张图储存了进行back propagation需要的中间结果(中间结果即每个需要更新变量的输出对本变量的局部梯度,用于根据链式法则进行梯度回传)。当调用了.backward()后,会从内存中将这张图进行释放。上述代码执行到************************时,内存中是包含了两张计算图的。执行到++++++时,得到对应的grad值并且释放内存。这样,训练时必须存储两张计算图,而如果loss的来源组成更加复杂,内存消耗会更大。
为了减小每次的内存消耗,借助梯度累加,有:
因此有如下变种
for idx, data in enumerate(train_loader):
xs, ys = data
optmizer.zero_grad()
# 计算d(l1)/d(x)
pred1 = model1(xs) #生成graph1
loss1 = loss_fn1(pred1, ys)
loss1.backward() #释放graph1
# 计算d(l2)/d(x)
pred2 = model2(xs)#生成graph2
loss2 = loss_fn2(pred2, ys)
loss2.backward() #释放graph2
# 使用d(l1)/d(x)+d(l2)/d(x)进行优化
optmizer.step()
由于梯度是默认累加的。因此即使在第一个模型forward结束,调用backward计算梯度并释放了计算图。第二个模型backward产生的梯度也能累加到第一个模型的梯度上。因此同一时刻,系统能始终保持最多存在一张计算图。借此可以减小内存的消耗。可以从代码中看出,利用梯度累加,可以在最多保存一张计算图的情况下进行multi-task任务的训练。
2、在内存大小不够的情况下,叠加多个batch的grad作为一个大batch进行迭代
传统的训练函数,一个batch是这么训练的:
for i,(images,target) in enumerate(train_loader):
# 1. input output
images = images.cuda(non_blocking=True)
target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
outputs = model(images)
loss = criterion(outputs,target)
# 2. backward
optimizer.zero_grad() # reset gradient
loss.backward()
optimizer.step()
- 获取loss:输入图像和标签,通过infer计算得到预测值,计算损失函数;
- optimizer.zero_grad() 清空过往梯度;
- loss.backward() 反向传播,计算当前梯度;
- optimizer.step() 根据梯度更新网络参数
简单的说就是进来一个batch的数据,计算一次梯度,更新一次网络
使用梯度累加是这么写的:
for i,(images,target) in enumerate(train_loader):
# 1. input output
images = images.cuda(non_blocking=True)
target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
outputs = model(images)
loss = criterion(outputs,target)
# 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
- 获取loss:输入图像和标签,通过infer计算得到预测值,计算损失函数;
- loss.backward() 反向传播,计算当前梯度;
- 多次循环步骤1-2,不清空梯度,使梯度累加在已有梯度上;
- 梯度累加了一定次数后,先optimizer.step() 根据累计的梯度更新网络参数,然后optimizer.zero_grad() 清空过往梯度,为下一波梯度累加做准备;
总结来说:梯度累加就是,每次获取1个batch的数据,计算1次梯度,梯度不清空,不断累加,累加一定次数后,根据累加的梯度更新网络参数,然后清空梯度,进行下一次循环。一定条件下,batchsize越大训练效果越好,梯度累加则实现了batchsize的变相扩大,如果accumulation_steps为8,则batchsize ‘变相’ 扩大了8倍,使用时需要注意,学习率也要适当放大(因为使用的样本增多,梯度更加稳定)。
有人会问,在上面的代码中为什么不直接对多个batch的loss先求和取平均,再梯度回传与更新。按我的理解这是为了减小内存的消耗。以accumulation_steps =8为例,采用多个batch的loss求和平均后再回传的方式,那么前提是对accumulation_steps 个batch进行了forward的过程,而每一次forward后都会生成一个计算图。也就是说,在这种方式下,生成了accumulation_steps 个计算图。 而采用上述代码的方式时,每次每个batch的forward结束,就进行backward计算梯度,backward结束也就释放了计算图,由于梯度是累加的,所以两种方法的计算结果是一致的。但是第二种方法在每一时刻最多只产生一张计算图,因此减小了内存消耗。