pytorch中的register_buffer

复制代码
今天在一个模型的init中遇到了self.register_buffer(‘running_mean’, torch.zeros(num_features))

register_buffer(self, name, tensor)是一个PyTorch中的方法,它的作用是向模块(module)中添加一个持久的缓冲区(buffer)

缓冲区是一种不被视为模型参数(model parameter)的张量(tensor),它不会在训练过程中更新梯度(gradient),但是会作为模块的状态(state)被保存和迁移

通常,缓冲区用于存储一些与模型相关但不需要学习的量,例如BatchNorm层的running_mean和running_var,它们是根据数据的统计信息计算得到的

缓冲区可以通过给定的名称作为属性(attribute)被访问,例如self.running_mean

参数name是一个字符串,表示缓冲区的名称。通过这个名称,可以从模块中获取缓冲区。

参数tensor是一个张量,表示缓冲区的初始值。当模块被创建时,缓冲区会被赋值为这个张量。

示例中,使用self.register_buffer('running_mean', torch.zeros(num_features))向模块中添加了一个名为running_mean的缓冲区,它的初始值是一个全零的张量,维度为num_features

缓冲区的创建和销毁是在什么时候?

相关推荐
金伟API10245 小时前
如何从零打造一个极简的DeepSeek-R1大模型
人工智能·ai
江畔柳前堤5 小时前
大语言模型分布式训练:从并行策略到万卡工程的系统梳理
人工智能·分布式·深度学习·算法·目标检测·机器学习·语言模型
Shockang5 小时前
智能体环路工程实战
人工智能
美摄科技5 小时前
美摄美颜特效SDK Skill:以AI赋能智能视音频新视界
人工智能
Web3_Daisy6 小时前
Pump.fun 与 FOMO 竞争背后的 Meme 市场变局
大数据·人工智能·金融·web3·区块链
AI创界者6 小时前
MiniMax-H3 本地一键部署整合包:8G 显存玩转文图生视频、视频参考、角色替换与超分补帧全流程
人工智能·深度学习
江畔柳前堤6 小时前
HBM:大语言模型时代的「算力血液」——从内存墙到带宽革命的深度拆解
服务器·人工智能·windows·目标检测·语言模型·自然语言处理·软件工程
美摄科技6 小时前
视频一键成片SDK Skill:AI智能分析与语义理解
人工智能
北斗落凡尘7 小时前
LangGraph 入门实战(2)
python·langchain
fthux7 小时前
装闭 RenoPit 源码解析(05):FastAPI与Celery如何执行AI装修分析
人工智能·ai·开源·github·open source·renopit