pytorch 加载使用部分预训练模型(pretrained model)

找了一些资料，对我这种菜鸟并不友好，把自己摸索的相对详细的过程记录一下。

pytorch 加载全部模型比较简单，直接使用如下代码:

net.load_state_dict(torch.load(pth_path))

现在只想使用 上述net（假设叫net-a）的中间的一部分模型，步骤如下：
1. 根据net-a的网络模型代码（如下），新建一个副本网络模型net-b代码。

    class Net-b(nn.Module):
        def __init__(self, **kwargs):
            super(Net-b, self).__init__()
            ... ... #保留net-a中 我们需要使用的中间卷积层等
        def forward():
        	... ... 

2. 根据Net-b网络模型文件，建立net-b

net-b = Net-b()  #别忘记传递必要的参数
net-b_dict = net-b.state_dict()

3. 将训练好的net-a权重 ‘复制’ 到net-b上

state_dict = torch.load(net-a_ckpt_path)	#加载预先训练好net-a的.pth文件
new_state_dict = OrderedDict()		#不是必要的【from collections import OrderedDict】 

new_state_dict = {k:v for k,v in state_dict.items() if k in net-b_dict}	#删除net-b不需要的键
net-b_dict.update(new_state_dict)	#更新参数
net-b.load_state_dict(net-b_dict)	#加载参数

检查一下是否复制成功：

for name, para in net-a.named_parameters():
    print(name, torch.max(para))
for name, para in net-b.named_parameters():
    print(name, torch.max(para))

打印部分结果如下，可见参数的名称和其中最大值是相同的，故net-b可以正常使用

#net-a 的参数和权重
b3.c3.body.0.bias tensor(0.0591, grad_fn=)
c1.body.0.weight tensor(0.3648, grad_fn=)
c1.body.0.bias tensor(0.0897, grad_fn=)
#net-b 的参数和权重
b3.c3.body.0.bias tensor(0.0591, grad_fn=)
c1.body.0.weight tensor(0.3648, grad_fn=)
c1.body.0.bias tensor(0.0897, grad_fn=)

除了这种方法外，还可以建立网络建构时，将中间使用块单独命名（net2），使用时直接调用Net.net2