pytorch只导入部分模型参数

我们在做迁移学习,或者在分割,检测等任务想使用预训练好的模型,同时又有自己修改之后的结构,使得模型文件保存的参数,有一部分是不需要的(don’t expected)。我们搭建的网络对保存文件来说,有一部分参数也是没有的(missed)。如果依旧使用torch.load(model.state_dict())的办法,就会出现 xxx expected,xxx missed类似的错误。那么在这种情况下,该如何导入模型呢?

好在Pytorch中的模型参数使用字典保存的,键是参数的名称,值是参数的具体数值。我们使用model.state_dict()获得这个字典,之后就能利用参数名称来实现导入。

请看下面的一个例子。

我们先搭建一个小小的网络。


   
     
     
     
     
  1. import torch as t
  2. from torch.nn import Module
  3. from torch import nn
  4. from torch.nn import functional as F
  5. class Net(Module):
  6. def __init__(self):
  7. super(Net,self).__init__()
  8. self.conv1 = nn.Conv2d( 3, 32, 3, 1)
  9. self.conv2 = nn.Conv2d( 32, 3, 3, 1)
  10. self.w = nn.Parameter(t.randn( 3, 10))
  11. for p in self.children():
  12. nn.init.xavier_normal_(p.weight.data)
  13. nn.init.constant_(p.bias.data, 0)
  14. def forward(self, x):
  15. out = self.conv1(x)
  16. out = self.conv2(x)
  17. out = F.avg_pool2d(out,(out.shape[ 2],out.shape[ 3]))
  18. out = F.linear(out,weight=self.w)
  19. return out

然后我们保存这个网络的初始值。


   
     
     
     
     
  1. model = Net()
  2. t.save(model.state_dict(), 'xxx.pth')

现在我们将Net修改一下,多加几个卷积层,但并不加入到forward中,仅仅出于少些几行的目的。


   
     
     
     
     
  1. import torch as t
  2. from torch.nn import Module
  3. from torch import nn
  4. from torch.nn import functional as F
  5. class Net(Module):
  6. def __init__(self):
  7. super(Net, self).__init__()
  8. self.conv1 = nn.Conv2d( 3, 32, 3, 1)
  9. self.conv2 = nn.Conv2d( 32, 3, 3, 1)
  10. self.conv3 = nn.Conv2d( 3, 64, 3, 1)
  11. self.conv4 = nn.Conv2d( 64, 32, 3, 1)
  12. for p in self.children():
  13. nn.init.xavier_normal_(p.weight.data)
  14. nn.init.constant_(p.bias.data, 0)
  15. self.w = nn.Parameter(t.randn( 3, 10))
  16. def forward(self, x):
  17. out = self.conv1(x)
  18. out = self.conv2(x)
  19. out = F.avg_pool2d(out, (out.shape[ 2], out.shape[ 3]))
  20. out = F.linear(out, weight=self.w)
  21. return out

 

我们现在试着导入之前保存的模型参数。


   
     
     
     
     
  1. path = 'xxx.pth'
  2. model = Net()
  3. model.load_state_dict(t.load(path))
  4. '''
  5. RuntimeError: Error(s) in loading state_dict for Net:
  6. Missing key(s) in state_dict: "conv3.weight", "conv3.bias", "conv4.weight", "conv4.bias".
  7. '''

出现了没有在模型文件中找到error中的关键字的错误。

现在我们这样导入模型


   
     
     
     
     
  1. path = 'xxx.pth'
  2. model = Net()
  3. save_model = t.load(path)
  4. model_dict = model.state_dict()
  5. state_dict = {k:v for k,v in save_model.items() if k in model_dict.keys()}
  6. print(state_dict.keys()) # dict_keys(['w', 'conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias'])
  7. model_dict.update(state_dict)
  8. model.load_state_dict(model_dict)

看看上面的代码,很容易弄明白。其中model_dict.update的作用是更新代码中搭建的模型参数字典。为啥更新我其实并不清楚,但这一步骤是必须的,否则还会报错。

为了弄清楚为什么要更新model_dict,我们不妨分别输出state_dict和model_dict的关键值看一看。


   
     
     
     
     
  1. for k in state_dict.keys():
  2. print(k)
  3. '''
  4. w
  5. conv1.weight
  6. conv1.bias
  7. conv2.weight
  8. conv2.bias
  9. '''

   
     
     
     
     
  1. for k in model_dict.keys():
  2. print(k)
  3. '''
  4. w
  5. conv1.weight
  6. conv1.bias
  7. conv2.weight
  8. conv2.bias
  9. conv3.weight
  10. conv3.bias
  11. conv4.weight
  12. conv4.bias
  13. '''

这个结果也是预料之中的,所以我猜测,update之后,model_dict和state_dict中具有相同键的值已经同步了。updata的目的就是使model_dict带有state_dict中都具有的那一部分参数的值,对于model_dict中有的,但是save_dict中没有的参数,值不改变,参数仍然使用初始值。

你可能感兴趣的:(学习工具及问题记录)