Pytorch预训练模型和修改——记录_pytorch加载预训练模型并更改维度

加载模型

一般从torchvision的models中加载常用模型，如alexnet、densenet、inception、resnet、squeezenet、vgg等常用网络结构，并提供预训练模型，调用方便。

from torchvision import models

resnet = models.resnet50(pretrain=True)
print(resnet)  # 打印网络结构
1
2
3
4

读取预训练模型

另一种是读取自己预训练模型，而不是使用官方自带。

import torch
resnet18 = models.resnet18(pretrained=False) #pretrained参数默认是False,为了代码清晰，最好还是加上参数赋值.

resnet18.load_state_dict(torch.load(path_params.pkl))
1
2
3
4

load_state_dict方法还有一个重要的参数是strict，该参数默认是True，表示预训练模型的层和自己定义的网络结构层严格对应相等（比如层名和维度）。

当新定义的网络（model_dict）和预训练网络（pretrained_dict）的层名不严格相等时，需要先将pretrained_dict里不属于model_dict的键剔除掉 ：

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 
1

再用预训练模型参数更新model_dict，最后用load_state_dict方法初始化自己定义的新网络结构。

整体代码：

print resnet18 #打印的还是网络结构
 
# 注意: cnn = resnet18.load_state_dict(torch.load( path_params.pkl )) #是错误的，这样cnn将是nonetype
 
pre_dict = resnet18.state_dict() #按键值对将模型参数加载到pre_dict
 
print for k, v in pre_dict.items(): # 打印模型参数
 
for k, v in pre_dict.items():
　　print k  #打印模型每层命名
 
# model是自己定义好的新网络模型，将pretrained_dict和model_dict中命名一致的层加入
 
# pretrained_dict（包括参数)。
 
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

模型参数修改

对模型特定层进行修改，一般直接调用对应层名并赋予新的层结构，常用是修改全连接层输出类别或特征数。

import torchvision.models as models 
#调用模型 
model = models.resnet50(pretrained=True) 
#提取fc层中固定的参数 
fc_features = model.fc.in_features 
#修改类别为9 
model.fc = nn.Linear(fc_features, 9)
1
2
3
4
5
6
7

训练特定层（冻结层）

对于自己定义的网络结构，需要选择特定层进行冻结时，往往需要用到 requires_grad函数来定义是否计算梯度。

count = 0
para_optim = []
for k in model.children():
　　count += 1
　　if count > 6:   # 6 should be changed properly
　　　　for param in k.parameters():
　　　　　　para_optim.append(param)
　　else:
　　　　for param in k.parameters():
　　　　　　param.requires_grad = False
optimizer = optim.RMSprop(para_optim, lr)#只对特定的层的参数进行优化更新，即选择特定的层进行finetune。
1
2
3
4
5
6
7
8
9
10
11

此代码实现了PyTorch中使用预训练的模型初始化网络的一部分参数，主要是：

• 设置选择条件
• 使用children函数得到子模型名，使用parameters得到参数
• 使用requires_grad确定不计算梯度，冻结层

PyTorch的Module.modules()和Module.children()

在PyTorch中，所有的neural network module都是class torch.nn.Module的子类，在Modules中可以包含其它的Modules，以一种树状结构进行嵌套。当需要返回神经网络中的各个模块时，Module.modules()方法返回网络中所有模块的一个iterator，而Module.children()方法返回所有直接子模块的一个iterator。具体而言：

list ( nn.Sequential(nn.Linear(10, 20), nn.ReLU()).modules() )
Out[9]:
[Sequential (
(0): Linear (10 -> 20)
(1): ReLU ()
), Linear (10 -> 20), ReLU ()]

In [10]: list( nn.Sequential(nn.Linear(10, 20), nn.ReLU()) .children() )
Out[10]: [Linear (10 -> 20), ReLU ()]
1
2
3
4
5
6
7
8
9

另一个例子：

下载好的模型，可以用下面这段代码看一下模型参数，并且改一下模型。在vgg19.pth同级目录建立一个test.py。

import torch
import torch.nn as nn
import torchvision.models as models

vgg16 = models.vgg16(pretrained=False)

#打印出预训练模型的参数
vgg16.load_state_dict(torch.load('vgg16-397923af.pth'))
print('vgg16:\n', vgg16) 

modified_features = nn.Sequential(*list(vgg16.features.children())[:-1])
# to relu5_3
print('modified_features:\n', modified_features )#打印修改后的模型参数

修改好之后features就可以拿去做Faster-RCNN提取特征用了
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

参考资料