模型保存

在 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 ，它以键值对的形式包含模型中需要保存下来的参数，例如:

classMyModule(nn.Module):def__init__(self, input_size, output_size):super(MyModule, self).__init__()
        self.lin= nn.Linear(input_size, output_size)defforward(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