pytorch梯度累积

Billie使劲学2024-03-03 22:16

梯度累加其实是为了变相扩大batch_size，用来解决显存受限问题。

常规训练方式，每次从train_loader读取出一个batch的数据：

python 复制代码

for x,y in train_loader:
	pred = model(x)
	loss = criterion(pred, label)
	# 反向传播
	loss.backward()
	# 根据新的梯度更新网络参数
	optimizer.step()
	# 清空以往梯度，通过下面反向传播重新计算梯度
	optimizer.zero_grad()

pytorch每次forward完都会得到一个用于梯度回传的计算图，pytorch构建的计算图是动态的，其实在每次backward后计算图都会从内存中释放掉，但是梯度不会清空的。所以若不显示的进行optimizer.zero_grad()清空过往梯度这一步操作，backward()的时候就会累加过往梯度。

梯度累加的做法：

python 复制代码

accumulation_steps = 4
for i,(x,y) in enumerate(train_loader):
	pred = model(x)
	loss = criterion(pred, label)
	
	# 相当于对累加后的梯度取平均
	loss = loss/accumulation_steps
	# 反向传播
	loss.backward()

	if (i+1) % accumulation_steps == 0:
		# 根据新的梯度更新网络参数
		optimizer.step()
		# 清空以往梯度，通过下面反向传播重新计算梯度
		optimizer.zero_grad()

代码中设置accumulation_steps = 4，意思就是变相扩大batch_size四倍。因为代码中每隔4次迭代才清空梯度，更新参数。

至于为啥loss = loss/accumulation_steps，因为梯度累加了四次呀，那就要取平均，除以4。那我每次loss取4，其实就相当于最后将累加后的梯度除4咯。同时，因为累计了4个batch，那学习率也应该扩大4倍，让更新的步子跨大点。

看网上的帖子有讨论对BN层是否有影响，因为BN的估算阶段（计算batch内均值、方差）是在forward阶段完成的，那真实的batch_size放大4倍效果肯定是比通过梯度累加放大4倍效果好的，毕竟计算真实的大batch_size内的均值、方差肯定更精确。

还有讨论说通过调低BN参数momentum可以得到更长序列的统计信息，应该意思是能够记忆更久远的统计信息（均值、方差），以逼近真实的扩大batch_size的效果。

参考

pytorch骚操作之梯度累加，变相增大batch size