pytorch 预训练模型冻结层+添加层+不同层学习率调整——以resnet50为例

前言

网上完整的教程比较少,很多讲的都是局部操作,比如:如何冻结层、如何添加层、如何调整不同层的学习率,一旦组合起来,就会发现总是有bug。我在这篇文章里,尽可能地把遇到的bug写一下。

步骤

先定义一个新的类,这个类用来构造自己需要的模型。

"""
参数:
num_classes:自己的分类任务要分的类别数

代码构造:
BackBone:pytorch官方的预训练模型
add_block:需要添加的fc层

"""
class Net(nn.Module):
    def __init__(self , num_classes = 14): 
        super(Net, self).__init__() 
        BackBone = torchvision.models.__dict__['resnet50'](pretrained=True)
#         BackBone = nn.Sequential(*list(model.children())[:-1]) # 这里有个坑!
        add_block = []
        add_block += [nn.Linear(1000, 512)]
        add_block += [nn.ReLU(True)]
        add_block += [nn.Dropout(0.15)]
        add_block += [nn.Linear(512, 128)]
        add_block += [nn.ReLU(True)]
        add_block += [nn.Linear(128, num_classes)]
        add_block = nn.Sequential(*add_block)
        
        self.BackBone = BackBone
        self.add_block = add_block  

    def forward(self, x):
        
        x = self.BackBone(x)
        x = self.add_block(x)
        
        return x

"""定义一个实例,名叫model_new"""
model_new = Net()

查看model_new的网络架构:

for name, value in model_new.named_parameters():
    print(name)

"""

输出:
BackBone.conv1.weight
BackBone.bn1.weight
BackBone.bn1.bias
BackBone.layer1.0.conv1.weight  # 注意这里,是BackBone.layer1
BackBone.layer1.0.bn1.weight
BackBone.layer1.0.bn1.bias
BackBone.layer1.0.conv2.weight
BackBone.layer1.0.bn2.weight
BackBone.layer1.0.bn2.bias
···

"""

注意:经过上面的操作后,网络结构会有一些微妙的变化,如下:

# 对原模型
for name, child in model.named_children():
    print(name)

"""
输出:
conv1
bn1
relu
maxpool
layer1
layer2
layer3
layer4
avgpool
fc
"""

# 对修改后的模型
for name, child in model_new.named_children():
    print(name)

"""
输出:
BackBone
add_block
"""

因此,对于model_new,需要进一步查看:(继续用 . 操作,进入到BackBone里)

for name, child in model_new.BackBone.named_children():
    print(name)

"""
输出:
conv1
bn1
relu
maxpool
layer1
layer2
layer3
layer4
avgpool
fc
"""

坑1:

网上比较流行的做法是,用这段代码来修改网络:

model = torchvision.models.__dict__['resnet50'](pretrained=True)
BackBone = nn.Sequential(*list(model.children())[:-1])

"""

输出:
BackBone.0.weight
BackBone.1.weight
BackBone.1.bias
BackBone.4.0.conv1.weight  
BackBone.4.0.bn1.weight
BackBone.4.0.bn1.bias
BackBone.4.0.conv2.weight
BackBone.4.0.bn2.weight
BackBone.4.0.bn2.bias
···

"""

这里是BackBone.4,这个.4就很迷,因为后面设置不同层学习率的话,需要用到
model_new.BackBone.谁.parameters(),问题出在这个“谁”身上,如果是layer4就没问题,但如果是4, model_new.BackBone.4就不符合Python语法规范,如下图:

(?)如果真要用这个方法,那么应该怎么定位到BackBone.4这里呢?求解答。 


 接着,冻结层操作,我需要冻结网络的前3个模块,只训练后面的模块:

"""
前面用model_new.BackBone.named_children()查看了有哪些“总的层”,那么就需要把这个当做标签。
比如,我只对最后一个模块(conv5),以及新添加的add_block做训练,那么就设置这些层的requires_grad = True
"""

for name, child in model_new.named_children(): 
    if name in ['add_block']: 
        for param in child.parameters():
            param.requires_grad = True
            
for name, child in model_new.BackBone.named_children():
    if name in ['layer4','avgpool','fc']:
        for param in child.parameters():
            param.requires_grad = True
    else:
        for param in child.parameters():
            param.requires_grad = False
print("模型冻结完成")

注意:.named_children()会忽略一些层, 比如最后的avgpool层,那么,我们可以直接看model_new长什么样子,然后把参数添加到列表里,如下图: pytorch 预训练模型冻结层+添加层+不同层学习率调整——以resnet50为例_第1张图片

 经过冻结操作后,查看一下所有层的冻结情况:

for name, value in model_new.named_parameters():
    print(name, "\t冻结=\t",value.requires_grad)

"""
输出:
BackBone.conv1.weight 	冻结=	 False
BackBone.bn1.weight 	冻结=	 False
BackBone.bn1.bias 	冻结=	 False
BackBone.layer1.0.conv1.weight 	冻结=	 False
···
add_block.3.weight 	冻结=	 True
add_block.3.bias 	冻结=	 True
add_block.5.weight 	冻结=	 True
add_block.5.bias 	冻结=	 True

"""

接下来,设置不同层的学习率:(一般来说,进行微调的层的设置的小一些,新加入的层的学习率设置的大一些,比如后者是前者的10-100倍)

"""
这里有不少坑,见后面分析

"""

optimizer = torch.optim.SGD([                         
                            {'params': model.module.BackBone.fc.parameters(),'lr': 0.1},
                            {'params': model.module.add_block.parameters(),'lr': 0.1},
                            {'params': model.module.BackBone.layer4.parameters(),'lr': 0.01},
                        ], momentum=0.9, weight_decay=1e-4)

坑2:如果用多GPU训练,会用到model=torch.nn.DataParallel(model)这句代码,那么需要把model.BackBone....改为model.module.BackBone....


 坑3:不能乱加 .parameters(),不然会报错。

网上关于如何冻结层、如何设置不同学习率的讲解,有不少用到这句代码:

filter(lambda p: ...... model_new.parameters())

如果把两个filter放到优化器里,那么就会得到这样的报错:

TypeError: __init__() got multiple values for argument 'lr'

这个错误的意思是,你定义了两个model_new.parameters(),但优化器只能接收一个。

例如:

""" 代码块1 """
for param in model_new.parameters():
    param.requires_grad = False
for param in model_new.add_block.parameters():
    param.requires_grad = True

""" 代码块2 """
ignored_params = list(map(id, model_new.add_block.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, model_new.parameters())


optimizer = optim.SGD(

            # 对应代码块1
            filter(lambda p: p.requires_grad, model_new.parameters()), 
            
            # 对应代码块2
            [
                {'params': base_params, 'lr': 0.01},
                {'params': model_new.add_block.parameters(), 'lr': 0.1},
            ], 

            weight_decay=1e-5, momentum=0.9)

两个红框框都包含了model_new.parameters()了,所以就重复了。

pytorch 预训练模型冻结层+添加层+不同层学习率调整——以resnet50为例_第2张图片

简而言之,在没有找到合适解决办法之前,冻结层的操作不需要用 filter(...),设置学习率的时候用filter(...)就行。


 最后,可以训练了。

你可能感兴趣的:(机器学习/深度学习,pytorch,python,深度学习,人工智能,神经网络)