pytorch保存和加载模型以及如何load部分参数

本文总结了pytorch中保存和加载模型的方法，以及在保存的模型文件与新定义的模型的参数不一一对应时，我们该如何加载模型参数。

1. 模型保存和加载的基本方式

在PyTorch中，模型可以通过两种方式保存和加载：保存整个模型（包括模型架构和参数）或仅保存模型的参数（state_dict）。

• 保存整个模型: 保存模型的架构和所有的权重参数。这样做的好处是可以直接加载使用，无需再定义模型架构，但是无法再对模型做出调整，不够灵活。

  python
  import torch
  import torchvision.models as models
  # 实例化一个预训练的resnet模型
  model = models.resnet18(pretrained=True)
  # 保存整个模型
  torch.save(model, 'model.pth')
  

  # 加载整个模型
  model = torch.load('model.pth')
  

• 仅保存模型参数
  通常推荐此方式，因为它仅保存权重参数，体积更小，更灵活，需要时可用新定义的模型结构加载参数。
  保存的参数通过model.state_dict()获取，得到一个有序字典类型：collections.OrderedDict，其中key是参数名称，value是保存了参数数值的tensor类型。

  OrderedDict是 Python 标准库 collections 模块中的一种字典（dict）类的子类。和普通的字典相比，OrderedDict 缴存了元素插入的顺序，所以当对其进行迭代时，键值对会按照添加的先后次序返回，而不是基于键的散列值。

  保存模型参数示例：

  # 保存模型的state_dict
  torch.save(model.state_dict(), 'model_state_dict.pth')
  

  加载模型参数示例：

  # 首先需要重新定义模型的结构，这里假设我们已经有了一模一样的模型定义
  model = models.resnet18(pretrained=False) # 取消预训练权重
  # 加载模型参数
  model.load_state_dict(torch.load('model_state_dict.pth'))
  

2. 保存的模型文件和当前定义的模型参数不完全一致时

有时候我们会对一个pretrained model的若干层进行一些修改，涉及到层的添加和减少，同时未改变的那些层想要load pretrained model的参数。
假设新定义的模型是new_net， pretrained模型是old_net, 以下两种方式适用于以下所有场景：
1. old_net的参数是new_net的子集
2. new_net的参数是old_net的子集
3. new_net和old_net的参数有交集

• strict=False
  一个直接的方式是在load_state_dict时strict=False，这样在load参数时pytorch会匹配两个模型中参数名字相同的参数进行导入。

  net_2.load_state_dict(torch.load("net_1.pth"), strict=False)
  

• 一种更灵活的方式，可自行添加更多的规则

  def load(save_path, model):
      pretraind_dict = torch.load(save_path)
      model_dict =  model.state_dict()
      # 只将pretraind_dict中那些在model_dict中的参数，提取出来
      state_dict = {k:v for k,v in pretraind_dict.items() if k in model_dict.keys()}
      # 将提取出来的参数更新到model_dict中，而model_dict有的，而state_dict没有的参数，不会被更新
      model_dict.update(state_dict)
      model.load_state_dict(model_dict)
  

  可利用上面的代码自行设计一些规则，比如如果不要laod某个参数，就可以在上面的代码中修改：

  state_dict = {k:v for k,v in pretraind_dict.items() if k in model_dict.keys() and k != 'conv1.weight'}
  

3. 验证代码

import torch
from torch import nn as nn

class model_2_convs(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 32, 3)
        self.mlp = nn.Linear(32, 10)
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        return x
    
class model_3_convs(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 32, 3)
        self.conv3 = nn.Conv2d(32, 64, 3)
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        return x
    


def load(save_path, model):
    pretraind_dict = torch.load(save_path)
    model_dict =  model.state_dict()
    # 只将pretraind_dict中那些在model_dict中的参数，提取出来
    state_dict = {k:v for k,v in pretraind_dict.items() if k in model_dict.keys()}
    # print(state_dict.keys())
    # 将提取出来的参数更新到model_dict中，而model_dict有的，而state_dict没有的参数，不会被更新
    model_dict.update(state_dict)
    model.load_state_dict(model_dict)

def load_weight_from_3_conv_to_2_conv(use_strict=False):
    net_1 = model_3_convs()
    net_2 = model_2_convs()

    torch.save(net_1.state_dict(), "net_1.pth")

    if use_strict:
        net_2.load_state_dict(torch.load("net_1.pth"), strict=False)
    else:
        load("net_1.pth", net_2)
    for key, para in net_2.state_dict().items():
        print(key)
        if key in net_1.state_dict().keys():
            print(torch.equal(para, net_1.state_dict()[key]))

def load_weight_from_2_conv_to_3_conv(use_strict=False):
    net_1 = model_3_convs()
    net_2 = model_2_convs()

    torch.save(net_2.state_dict(), "net_2.pth")
    if use_strict:
        net_1.load_state_dict(torch.load("net_2.pth"), strict=False)
    else:
        load("net_2.pth", net_1)
    for key, para in net_1.state_dict().items():
        print(key)
        if key in net_2.state_dict().keys():
            print(torch.equal(para, net_2.state_dict()[key]))
    



if __name__ == "__main__":
    load_weight_from_3_conv_to_2_conv(use_strict=True)
    load_weight_from_3_conv_to_2_conv(use_strict=False)
    load_weight_from_2_conv_to_3_conv(use_strict=True)
    load_weight_from_2_conv_to_3_conv(use_strict=False)