Pytorch模型中的parameter与buffer

模型保存

在 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
• 一种是反向传播不需要被optimizer更新，称之为 buffer

第一种参数我们可以通过 model.parameters() 返回；第二种参数我们可以通过 model.buffers() 返回。因为我们的模型保存的是 state_dict 返回的 OrderDict，所以这两种参数不仅要满足是否需要被更新的要求，还需要被保存到OrderDict。

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

• 第一种参数有两种方式:
  1.我们可以直接将模型的成员变量(self.xxx) 通过nn.Parameter() 创建，会自动注册到parameters中，可以通过model.parameters()返回，并且这样创建的参数会自动保存到OrderDict中去；
  2.通过nn.Parameter()创建普通Parameter对象，不作为模型的成员变量，然后将Parameter对象通过register_parameter()进行注册，可以通model.parameters()返回，注册后的参数也会自动保存到OrderDict中去;
• 第二种参数我们需要创建tensor, 然后将tensor通过register_buffer()进行注册，可以通model.buffers() 返回，注册完后参数也会自动保存到OrderDict中去。

分析一个例子：

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        buffer = torch.randn(2, 3)  # tensor
        self.register_buffer('my_buffer', buffer)
        self.param = nn.Parameter(torch.randn(3, 3))  # 模型的成员变量

    def forward(self, x):
        # 可以通过 self.param 和 self.my_buffer 访问
        pass
model = MyModel()
for param in model.parameters():
    print(param)
print("----------------")
for buffer in model.buffers():
    print(buffer)
print("----------------")
print(model.state_dict())

输出结果:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        buffer = torch.randn(2, 3)  # tensor
        param = nn.Parameter(torch.randn(3, 3))  # 普通 Parameter 对象
        self.register_buffer('my_buffer', buffer)
        self.register_parameter("param", param)

    def forward(self, x):
        # 可以通过 self.param 和 self.my_buffer 访问
        pass
model = MyModel()
for param in model.parameters():
    print(param)
print("----------------")
for buffer in model.buffers():
    print(buffer)
print("----------------")
print(model.state_dict())

输出:

register_buffer的函数原型:

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

register_parameter的函数原型:

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

为什么不直接将不需要进行参数修改的变量作为模型类的成员变量就好了，还要进行注册?
有两个原因:
1.不进行注册，参数不能保存到 OrderDict，也就无法进行保存
2.模型进行参数在CPU和GPU移动时, 执行model.to(device)，注册后的参数也会自动进行设备移动

看一个例子:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.my_tensor = torch.randn(1) # 参数直接作为模型类成员变量
        self.register_buffer('my_buffer', torch.randn(1)) # 参数注册为 buffer
        self.my_param = nn.Parameter(torch.randn(1))

    def forward(self, x):
            return x

model = MyModel()
print(model.state_dict())
model.cuda()
print(model.my_tensor)
print(model.my_buffer)

输出结果:

可以看到模型类的成员变量不在OrderDict中，不能进行保存；模型在进行设备移动时，模型类的成员变量没有进行移动。

总结

1.模型中需要进行更新的参数注册为Parameter，不需要进行更新的参数注册为buffer
2.模型保存的参数是 model.state_dict() 返回的 OrderDict
3.模型进行设备移动时，模型中注册的参数(Parameter和buffer)会同时进行移动

转载于https://zhuanlan.zhihu.com/p/89442276