pytorch梯度累积

梯度累加其实是为了变相扩大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

相关推荐
啥也不行就是菜3 分钟前
【AI助手】从零构建文章抓取器 MCP(Node.js 版)
人工智能·mcp·trae
亚里随笔6 分钟前
ReSpec:突破RL训练瓶颈的推测解码优化系统
人工智能·深度学习·自然语言处理·大语言模型·rlhf
Anson Jiang34 分钟前
PyTorch轻松实现CV模型:零基础到实战
pytorch·python·django·flask·python开发
腾讯云开发者36 分钟前
对话香港城市大学张泽松:AI时代教育“变天”?先抓核心能力|TVP专访
人工智能
岁月宁静1 小时前
图像生成接口的工程化设计与落地实践:封装豆包图像生成模型 Seedream 4.0 API
前端·人工智能·node.js
风雨同舟的代码笔记1 小时前
5.Python函数与模块化工程实战:构建高复用代码体系
python
万岳科技程序员小金1 小时前
多商户商城APP源码开发的未来方向:云原生、电商中台与智能客服
人工智能·云原生·开源·软件开发·app开发·多商户商城系统源码·多商户商城app开发
蓝色 - Lanse1 小时前
模型推理如何利用非前缀缓存
人工智能·缓存
CoookeCola1 小时前
MovieNet (paper) :推动电影理解研究的综合数据集与基准
数据库·论文阅读·人工智能·计算机视觉·视觉检测·database
我的xiaodoujiao1 小时前
使用 Python 语言 从 0 到 1 搭建完整 Web UI自动化测试学习系列 22--数据驱动--参数化处理 Json 文件
python·学习·测试工具·pytest