参考文章:
pytorch中的pre-train函数模型引用及修改(增减网络层,修改某层参数等)
(继)pytorch中的pretrain模型网络结构修改
综合来讲,有三种方法
一是加载完模型后直接修改(适用于修改参数)
#调用模型
model = models.resnet50(pretrained=True)
#提取fc层中固定的参数
fc_features = model.fc.in_features
#修改类别为9
model.fc = nn.Linear(fc_features, 9)
这种情况,适用于修改的层在原模型中有名字(可以由self.层的名字获取到)。
如果层在sequential中。因为sequential类型没有定义具体名称,所以只能用数组+索引序号的方式获取。例如model.seq是一个sequential,那么其中的子层就用model.seq[0]、model.seq[n]来获得。
二是修改网络结构(增减卷积层等)
第一种方法只适用于简单的参数修改,有的时候我们往往要修改网络中的层次结构,这时只能用参数覆盖的方法,即自己先定义一个类似的网络,再将预训练中的参数提取到自己的网络中来。
相见参考文献
三是使用nn.module的model.children()的函数,重新定义自己model的层。这个比较灵活高效。
那么对于网络的修改怎样可以快速的去除model本身的网络层呢?
一个继承nn.module的model它包含一个叫做children()的函数,这个函数可以用来提取出model每一层的网络结构,在此基础上进行修改即可,修改方法如下(去除后两层):
resnet_layer = nn.Sequential(*list(model.children())[:-2])
那么,接下来就可以构建我们的网络了:
- 第一步:
class Net(nn.Module):
def __init__(self , model): #此处的model参数是已经加载了预训练参数的模型,方便继承预训练成果
super(Net, self).__init__()
#取掉model的后两层
self.resnet_layer = nn.Sequential(*list(model.children())[:-2])
self.transion_layer = nn.ConvTranspose2d(2048, 2048, kernel_size=14, stride=3)
self.pool_layer = nn.MaxPool2d(32)
self.Linear_layer = nn.Linear(2048, 8)
def forward(self, x):
x = self.resnet_layer(x)
x = self.transion_layer(x)
x = self.pool_layer(x)
x = x.view(x.size(0), -1)
x = self.Linear_layer(x)
return x
第二步:加载原模型的预训练参数,然后继承原模型的部分结构及相应参数,构建新架构
resnet = models.resnet50(pretrained=True) # 原模型,并加载预训练参数
model = Net(resnet) # 继承原模型的部分结构及相应参数,构建新架构