找了一些资料,对我这种菜鸟并不友好,把自己摸索的相对详细的过程记录一下。
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=<MaxBackward1>)
c1.body.0.weight tensor(0.3648, grad_fn=<MaxBackward1>)
c1.body.0.bias tensor(0.0897, grad_fn=<MaxBackward1>)
#net-b 的参数和权重
b3.c3.body.0.bias tensor(0.0591, grad_fn=<MaxBackward1>)
c1.body.0.weight tensor(0.3648, grad_fn=<MaxBackward1>)
c1.body.0.bias tensor(0.0897, grad_fn=<MaxBackward1>)
除了这种方法外,还可以建立网络建构时,将中间使用块单独命名(net2),使用时直接调用Net.net2