Pytorch中实现只导入部分模型参数的方式
我们在做迁移学习，或者在分割，检测等任务想使用预训练好的模型，同时又有自己修改之后的结构，使得模型文件保存的参数，有一部分是不需要的（don'texpected）。我们搭建的网络对保存文件来说，有一部分参数也是没有的（missed）。如果依旧使用torch.load（model.state_dict())的办法，就会出现xxxexpected，xxxmissed类似的错误。那么在这种情况下，该如何导入模型呢？
好在Pytorch中的模型参数使用字典保存的，键是参数的名称，值是参数的具体数值。我们使用model.state_dict()获得这个字典，之后就能利用参数名称来实现导入。
请看下面的一个例子。
我们先搭建一个小小的网络。
importtorchast fromtorch.nnimportModule fromtorchimportnn fromtorch.nnimportfunctionalasF classNet(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)) forpinself.children(): nn.init.xavier_normal_(p.weight.data) nn.init.constant_(p.bias.data,0) defforward(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) returnout
然后我们保存这个网络的初始值。
model=Net() t.save(model.state_dict(),'xxx.pth')
现在我们将Net修改一下，多加几个卷积层，但并不加入到forward中，仅仅出于少些几行的目的。
importtorchast fromtorch.nnimportModule fromtorchimportnn fromtorch.nnimportfunctionalasF classNet(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) forpinself.children(): nn.init.xavier_normal_(p.weight.data) nn.init.constant_(p.bias.data,0) self.w=nn.Parameter(t.randn(3,10)) defforward(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) returnout
我们现在试着导入之前保存的模型参数。
path='xxx.pth' model=Net() model.load_state_dict(t.load(path)) ''' RuntimeError:Error(s)inloadingstate_dictforNet: Missingkey(s)instate_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:vfork,vinsave_model.items()ifkinmodel_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的关键值看一看。
forkinstate_dict.keys(): print(k) ''' w conv1.weight conv1.bias conv2.weight conv2.bias ''' forkinmodel_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中实现只导入部分模型参数的方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持毛票票。
声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。