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

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

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

# 保存模型的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'))

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)

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

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)