Pytorch parameter与buffer_model.buffer()中的参数-优快云博客

本文介绍了PyTorch中模型的保存与加载方法，重点讨论了Parameter和Buffer的区别。Parameter是需要被optimizer更新的，而Buffer则不需要。创建Parameter可以通过nn.Parameter()直接赋值给模型成员或使用register_parameter()注册。Buffer则需要通过register_buffer()注册。模型保存时，state_dict会包含这两类参数。加载时，通过load_state_dict()恢复模型状态。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

转载自知乎
原文链接：https://zhuanlan.zhihu.com/p/89442276

模型保存

在 Pytorch 中一种模型保存和加载的方式如下:

# save
torch.save(model.state_dict(), PATH)

# load
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

可以看到模型保存的是 model.state_dict() 的返回对象。 model.state_dict() 的返回对象是一个 OrderDict ，它以键值对的形式包含模型中需要保存下来的参数，例如:

class MyModule(nn.Module):
    def __init__(self, input_size, output_size):
        super(MyModule, self).__init__()
        self.lin = nn.Linear(input_size, output_size)
    def forward(self, x):
        return self.lin(x)

module = MyModule(4, 2)
print(module.state_dict())

在这里插入图片描述
模型中的参数就是线性层的 weight 和 bias.

Parameter和buffer

模型中需要保存下来的参数包括两种:
一种是反向传播需要被optimizer更新的，称之为 parameter。可以通过 model.parameters() 返回。
一种是反向传播不需要被optimizer更新，称之为 buffer。可以通过 model.buffers() 返回。

模型进行设备移动时，模型中注册的参数(Parameter和buffer)会同时进行移动。

因为我们的模型保存的是 state_dict 返回的 OrderDict，所以这两种参数不仅要满足是否需要被更新的要求，还需要被保存到OrderDict。

那么现在的问题是这两种参数如何创建呢，创建好了如何保存到OrderDict呢？

Parameter

有两种方式创建并保存到OrderDict:

1、可以直接将模型的成员变量(self.xxx) 通过nn.Parameter() 创建，会自动注册到parameters中，可以通过model.parameters() 返回，并且这样创建的参数会自动保存到OrderDict中去；

2、通过nn.Parameter() 创建普通Parameter对象，不作为模型的成员变量，然后将Parameter对象通过register_parameter()进行注册，可以通model.parameters() 返回，注册后的参数也会自动保存到OrderDict中去;

register_parameter的函数原型:

register_parameter(name, param) name: string param: Parameter

buffer

有一种方式创建并保存到OrderDict:

需要创建tensor, 然后将tensor通过register_buffer()进行注册，可以通model.buffers() 返回，注册完后参数也会自动保存到OrderDict中去。

register_buffer的函数原型:

register_buffer(name, tensor) name: string tensor: Tensor