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长什么样子，然后把参数添加到列表里，如下图： 

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

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()了，所以就重复了。

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

---

 最后，可以训练了。