快捷導(dǎo)航

Pytorch中實(shí)現(xiàn)只導(dǎo)入部分模型參數(shù)的方式

更新時(shí)間：2020年01月02日 16:57:09 作者：咆哮的阿杰

今天小編就為大家分享一篇Pytorch中實(shí)現(xiàn)只導(dǎo)入部分模型參數(shù)的方式，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過來看看吧

我們?cè)谧鲞w移學(xué)習(xí)，或者在分割，檢測(cè)等任務(wù)想使用預(yù)訓(xùn)練好的模型，同時(shí)又有自己修改之后的結(jié)構(gòu)，使得模型文件保存的參數(shù)，有一部分是不需要的（don't expected）。我們搭建的網(wǎng)絡(luò)對(duì)保存文件來說，有一部分參數(shù)也是沒有的（missed）。如果依舊使用torch.load（model.state_dict())的辦法，就會(huì)出現(xiàn) xxx expected，xxx missed類似的錯(cuò)誤。那么在這種情況下，該如何導(dǎo)入模型呢？

好在Pytorch中的模型參數(shù)使用字典保存的，鍵是參數(shù)的名稱，值是參數(shù)的具體數(shù)值。我們使用model.state_dict()獲得這個(gè)字典，之后就能利用參數(shù)名稱來實(shí)現(xiàn)導(dǎo)入。

請(qǐng)看下面的一個(gè)例子。

我們先搭建一個(gè)小小的網(wǎng)絡(luò)。

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

然后我們保存這個(gè)網(wǎng)絡(luò)的初始值。

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

現(xiàn)在我們將Net修改一下，多加幾個(gè)卷積層，但并不加入到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

我們現(xiàn)在試著導(dǎo)入之前保存的模型參數(shù)。

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". 
'''

出現(xiàn)了沒有在模型文件中找到error中的關(guān)鍵字的錯(cuò)誤。

現(xiàn)在我們這樣導(dǎo)入模型

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的作用是更新代碼中搭建的模型參數(shù)字典。為啥更新我其實(shí)并不清楚，但這一步驟是必須的，否則還會(huì)報(bào)錯(cuò)。

為了弄清楚為什么要更新model_dict，我們不妨分別輸出state_dict和model_dict的關(guān)鍵值看一看。

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
'''

這個(gè)結(jié)果也是預(yù)料之中的，所以我猜測(cè)，update之后，model_dict和state_dict中具有相同鍵的值已經(jīng)同步了。updata的目的就是使model_dict帶有state_dict中都具有的那一部分參數(shù)的值，對(duì)于model_dict中有的，但是save_dict中沒有的參數(shù)，值不改變，參數(shù)仍然使用初始值。

以上這篇Pytorch中實(shí)現(xiàn)只導(dǎo)入部分模型參數(shù)的方式就是小編分享給大家的全部?jī)?nèi)容了，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: