第二十四篇—模型中的parameter和buffer
YouOnly_LiveOnce
编辑于 2023年04月01日 02:11
收录于文集
共28篇

部分一：模型保存

模型保存方式：

代码块
Python
自动换行
复制代码
# save
torch.save(model.state_dict(),PATH)
复制成功

模型加载方式：

代码块
Python
自动换行
复制代码
# load
model=Model(*args,**kwargs)
model.load_state_dict(torch.load(PATH))
复制成功

保存的内容是model.state_dict()的返回对象，是一个OrderedDict，以键值对(key-val)的形式包含模型中需要保存下来的参数，如：

代码块
Python
自动换行
复制代码
import torch.nn as nn
# 定义模型
class Model(nn.Module):
    def __init__(self,input,output):
        super(Model,self).__init__()
        self.layer=nn.Linear(input,output)
    def forward(self,x):
        return self.net(x)
# 创建模型
model=Model(input=4,output=1)
# 打印模型
print(model)
# 打印模型参数
print(model.state_dict())
复制成功

打印模型：

打印模型参数：

部分二：模型中的parameter和buffer

parameter在反向传播时会被optimizer.step更新，而buffer在反向传播时不会被更新；

parameter和buffer都保存在model.state_dict()返回的OrderedDict对象中；

模型进行设备移动时，模型中注册的参数(parameter和buffer)，即model.state_dict()中的内容会同时进行移动。

创建parameter和buffer

创建parameter：

方法一：直接将模型的成员变量self.xxx通过nn.Parameter()创建，会自动注册到model.parameters()中；

代码块
Python
自动换行
复制代码
# 方法一
self.param=nn.Parameter(torch.randn(4,1))
复制成功

方法二：先通过nn.Parameter()创建普通的parameter对象，此时该对象不作为模型的成员变量，然后将parameter对象通过register_parameter()注册到model.parameters()中；

代码块
Python
自动换行
复制代码
# 方法二
param=nn.Parameter(torch.randn(4,1))
self.register_parameter("param",param)
复制成功

创建buffer：

通过register_buffer()注册到model.buffers()中；

代码块
Python
自动换行
复制代码
self.register_buffer("buffer",torch.randn(3,1))
复制成功

注意一：parameter和buffer都保存在model.state_dict()中，举例如下：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn
# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        # 创建4*1的parameter(标准正态分布)
        self.param=nn.Parameter(torch.randn(4,1))
        # 创建2*1的buffer(标准正态分布)
        self.register_buffer("buffer",torch.randn(2,1))
    def forward(self,x):
        pass 
# 创建模型
model=Model()
# 打印模型参数
print(model.state_dict())
print("****************")
# 打印模型parameter
for param in model.parameters():
    print(param)
print("****************")
# 打印模型buffer
for buffer in model.buffers():
    print(buffer)
复制成功

结果如下：

注意二：需要进行梯度更新的模型参数属于parameter，而不是buffer中。实际上，parameter的创建在torch.nn.Linear类中的__init__函数中完成，成员变量weight和bias都属于parameter对象，并在这里进行了初始化，__init__函数如下图：

代码举例：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn
# 定义模型
class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        # 创建4*1的parameter(标准正态分布)
        # self.param=nn.Parameter(torch.randn(4,1))
        self.layer=nn.Linear(4,1)
        # 创建2*1的buffer(标准正态分布)
        self.register_buffer("buffer",torch.randn(2,1))
    def forward(self,x):
        return self.layer(x)
# 创建模型
model=Model()
# 打印模型
print(model)
# 打印模型参数(包括:parameter和buffer)
print(model.state_dict())
复制成功

打印结果：

parameter在反向传播时会被optimizer.step更新，而buffer在反向传播时不会被更新，举例如下：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn
# 定义模型
class Model(nn.Module):
    def __init__(self,num_input,num_output):
        super(Model,self).__init__()
        # 创建4*1的parameter(标准正态分布)
        self.layer=nn.Linear(num_input,num_output)
        # 创建2*1的buffer(标准正态分布)
        self.register_buffer("buffer",torch.randn(2,1))
    def forward(self,x):
        return self.layer(x)
# 创建模型
model=Model(4,1)
# 更新前参数打印
print("更新前:")
print(model.state_dict())
# 创建损失函数
loss_fn=nn.L1Loss()
# 创建优化器
optimizer=torch.optim.SGD(model.parameters(),lr=0.01)
# 创建一个4维输入
input=torch.ones(4,requires_grad=True)
# 前向传播
output=model(input)
# 设置标签
target=torch.ones([1])
# 计算损失
loss=loss_fn(output,target)
# 梯度清零
optimizer.zero_grad()
# 反向传播
loss.backward()
# 梯度更新
optimizer.step()
# 更新后参数打印
print("更新后:")
print(model.state_dict())
复制成功

结果如下(buffer不变；parameter更新)：

注意三：数据转换为buffer的原因

原因一：parameter和buffer都保存在model.state_dict()返回的OrderedDict对象中，如果保存模型的方法是torch.save(model.state_dict(),PATH)，那么类中普通成员数据如果不注册为parameter或buffer，就不能被保存，因为只有model.state_dict()中的内容会被保存，即只保存parameter和buffer；

原因二：模型进行设备移动时，模型中注册的参数(parameter和buffer)，即model.state_dict()中的内容会同时进行移动，而类中普通成员数据如果不注册为parameter或buffer，就不能被移动，如下所示：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn
# 定义模型
class Model(nn.Module):
    def __init__(self,num_input,num_output):
        super(Model,self).__init__()
        # 创建4*1的parameter(标准正态分布)
        self.layer=nn.Linear(num_input,num_output)
        # 创建2*1的buffer(标准正态分布)
        self.register_buffer("buffer",torch.randn(2,1))
        # 创建普通成员变量(标准正态分布)
        self.a=torch.randn(3)
    def forward(self,x):
        return self.layer(x)
# 创建模型
model=Model(4,1)
# 设备转移:CPU-GPU
model.cuda()
print(model.state_dict())
print("************")
print(model.a)
复制成功

结果如下(只有model.state_dict()中的内容，即parameter和buffer，转移到了GPU；而普通成员变量没被转移，仍在CPU上)：

深度学习 pytorch
本文为我原创，未经授权禁止转载
cv22768393