Pytorch中实现只导入部分模型参数的方式
我们在做迁移学习，或者在分割，检测等任务想使用预训练好的模型，同时又有自己修改之后的结构，使得模型文件保存的参数，有一部分是不需要的（don't expected）。我们搭建的网络对保存文件来说，有一部分参数也是没有的（missed）。如果依旧使用torch.load（model.state_dict())的办法，就会出现 xxx expected，xxx missed类似的错误。那么在这种情况下，该如何导入模型呢？
好在Pytorch中的模型参数使用字典保存的，键是参数的名称，值是参数的具体数值。我们使用model.state_dict()获得这个字典，之后就能利用参数名称来实现导入。
请看下面的一个例子。
我们先搭建一个小小的网络。
import torch as t
from torch.nn import Module
from torch import nn
from torch.nn import functional as F
class Net(Module):
def __init__(self):
super(Net,self).__init__()
self.conv1 = nn.Conv2d(3,32,3,1)
self.conv2 = nn.Conv2d(32,3,3,1)
self.w = nn.Parameter(t.randn(3,10))
for p in self.children():
nn.init.xavier_normal_(p.weight.data)
nn.init.constant_(p.bias.data, 0)
def forward(self, x):
out = self.conv1(x)
out = self.conv2(x)
out = F.avg_pool2d(out,(out.shape[2],out.shape[3]))
out = F.linear(out,weight=self.w)
return out
然后我们保存这个网络的初始值。
model = Net() t.save(model.state_dict(),'xxx.pth')
现在我们将Net修改一下，多加几个卷积层，但并不加入到forward中，仅仅出于少些几行的目的。
import torch as t
from torch.nn import Module
from torch import nn
from torch.nn import functional as F
class Net(Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 3, 3, 1)
self.conv3 = nn.Conv2d(3,64,3,1)
self.conv4 = nn.Conv2d(64,32,3,1)
for p in self.children():
nn.init.xavier_normal_(p.weight.data)
nn.init.constant_(p.bias.data, 0)
self.w = nn.Parameter(t.randn(3, 10))
def forward(self, x):
out = self.conv1(x)
out = self.conv2(x)
out = F.avg_pool2d(out, (out.shape[2], out.shape[3]))
out = F.linear(out, weight=self.w)
return out
我们现在试着导入之前保存的模型参数。
path = 'xxx.pth' model = Net() model.load_state_dict(t.load(path)) ''' RuntimeError: Error(s) in loading state_dict for Net: Missing key(s) in state_dict: "conv3.weight", "conv3.bias", "conv4.weight", "conv4.bias". '''
出现了没有在模型文件中找到error中的关键字的错误。
现在我们这样导入模型
path = 'xxx.pth'
model = Net()
save_model = t.load(path)
model_dict = model.state_dict()
state_dict = {k:v for k,v in save_model.items() if k in model_dict.keys()}
print(state_dict.keys()) # dict_keys(['w', 'conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias'])
model_dict.update(state_dict)
model.load_state_dict(model_dict)
看看上面的代码，很容易弄明白。其中model_dict.update的作用是更新代码中搭建的模型参数字典。为啥更新我其实并不清楚，但这一步骤是必须的，否则还会报错。
为了弄清楚为什么要更新model_dict，我们不妨分别输出state_dict和model_dict的关键值看一看。
for k in state_dict.keys(): print(k) ''' w conv1.weight conv1.bias conv2.weight conv2.bias ''' for k in model_dict.keys(): print(k) ''' w conv1.weight conv1.bias conv2.weight conv2.bias conv3.weight conv3.bias conv4.weight conv4.bias '''
这个结果也是预料之中的，所以我猜测，update之后，model_dict和state_dict中具有相同键的值已经同步了。updata的目的就是使model_dict带有state_dict中都具有的那一部分参数的值，对于model_dict中有的，但是save_dict中没有的参数，值不改变，参数仍然使用初始值。
以上这篇Pytorch中实现只导入部分模型参数的方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持脚本之家。
相关文章
这篇文章主要介绍了matplotlib绘制多子图共享鼠标光标的方法示例，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2021-01-01
这篇文章主要介绍了python检测服务器端口代码实例,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下2019-08-08
圣诞节快到了,分别用R和Python绘制了圣诞树祝你们圣诞节快乐,所以下面这篇文章主要给大家介绍了关于如何利用python/R绘制圣诞树的相关资料,需要的朋友可以参考下2021-12-12
这篇文章主要介绍了python之如何将标签转化为one-hot(独热编码)问题，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教2023-06-06
本篇文章主要介绍Python 的迭代器与zip,可迭代对象的相关概念,有需要的小伙伴可以参考下,希望能够给你带来帮助2021-11-11
这篇文章主要介绍了Python实现CET查分的方法,实例分析了Python操作链接查询的技巧,需要的朋友可以参考下2015-03-03
这篇文章主要为大家详细介绍了Python设计模式中的备忘录模式，文中示例代码介绍的非常详细，具有一定的参考价值，感兴趣的小伙伴们可以参考一下，希望能够给你带来帮助2022-02-02
这篇文章主要给大家介绍了十个让你相见恨晚的Python骚操作，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2020-11-11
这篇文章主要介绍了Python实现简单状态框架的方法,涉及Python状态框架的实现技巧,具有一定参考借鉴价值,需要的朋友可以参考下2015-03-03
这篇文章主要介绍了Python实时监控网站浏览记录实现过程详解,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下2020-07-07
最新评论