pytorch保存和加载模型的方法及如何load部分参数
作者：BigerBang
本文总结了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)到此这篇关于pytorch保存和加载模型以及如何load部分参数的文章就介绍到这了,更多相关pytorch保存和加载模型内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！