Pytorch中实现只导入部分模型参数的方式

浅浅的花香味﹌ 2024-03-22 20:42 89阅读 0赞

![e0d34bcb3f6a87a90eaa31d14d350423.jpeg][]

今天小编就为大家分享一篇Pytorch中实现只导入部分模型参数的方式，具有很好的参考价值，希望对大家有所帮助。一起跟随微点阅读小编过来看看吧

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

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

请看下面的一个例子。

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

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

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

<table> 
 <tbody> 
 <tr> 
 <td> 1 2 </td> 
 <td> <code>model </code><code>=</code> <code>Net()</code> <code>t.save(model.state_dict(),</code><code>'xxx.pth'</code><code>)</code> </td> 
 </tr> 
 </tbody> 
</table>

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

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

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

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

现在我们这样导入模型

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

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

为了弄清楚为什么要更新model\_dict，我们不妨分别输出state\_dict和model\_dict的关键值看一看。

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

以上这篇Pytorch中实现只导入部分模型参数的方式就是小编分享给大家的全部内容了，希望能给大家一个参考

来源：[微点阅读][Link 1] [https://www.weidianyuedu.com][Link 1]

[e0d34bcb3f6a87a90eaa31d14d350423.jpeg]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/14/c08a2f561ebf4a24b603551cddc5be3c.png
[Link 1]: https://www.weidianyuedu.com