Pytorch：找到对应的层，并在不同层设置不同的学习率

在使用一些预训练模型的过程中，通常要对“预训练模型”，“自己自定义的层”设置不同的学习率。
那么如何去设置不同的学习率呢？

0. 准备

为了可视化的方便，推荐使用pip install torch-summary安装 torch-summary 这个包，可以用来可视化网络。

使用指南见文档 torch-summary

1. 找到对应的层

pytorch的模型是基于类构建的，目前找到的最好方法是直接 print对应的模型，这里我们使用resnet50举例

import torchvision
model = torchvision.models.resnet50()
print(model) 

# print(model)的输出
# 我删掉了layer1，layer2，layer3，layer4，里面的结构为了防止看不完...
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(  # 这里太多了，显示不完了，建议大家自己看一下
      (0): Bottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
		# 没显示完
      )
      (1): Bottleneck(
      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
	# 没显示完
    )
  )
    (layer2): Sequential(  # 这里太多了，显示不完了，建议大家自己看一下
  )
    (layer3): Sequential(  # 这里太多了，显示不完了，建议大家自己看一下
  )
    (layer4): Sequential(  # 这里太多了，显示不完了，建议大家自己看一下
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)

这里我们可以看到每一行的的开始括号内都有对应的名字，我们便可以依据这些名字确定一些层，再让pytorch知道设置不同的学习率。

这里需要注意的是，有些同学可能会发现这些名字不是唯一的，比如说conv1这个名字就出现了三次。

但是实际上是不影响的，因为第二、三个conv1是包裹在Sequential内的，举个例子：

print(model.conv1)

# 下面是输出
# 下面是输出
Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)

对应参数我们可以发现model.conv1指向的就是第一个（第一层）的nn.Conv2d

那么如果我想获得第二个conv1呢（包裹在Sequential中的）？

可以注意到第二个conv1被包裹在了 (layer1): Sequential( (0)Bottleneck内

 (layer1): Sequential(  # 这里太多了，显示不完了，建议大家自己看一下
      (0): Bottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)

因此要迭代一下才能获得

print(model.layer1[0])
print('*'*50)
print(model.layer1[0].conv1)

# 下面是输出
# 下面是输出
Bottleneck(
  (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (downsample): Sequential(
    (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  )
)
**************************************************
Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)

可以看到我们就取出了对应的层，当然对应的block（比如这里的Bottleneck）其实也是这个方法取出来

2. 设置学习率

比如我们想对layer1的第0个Bottleneck设置学习率0.1，layer1的第1个Bottleneck中的conv1设置学习率0.2，其他的学习率为0.5

from torch import optim
block1 = list(map(id, model.layer1[0].parameters()))
conv1 = list(map(id, model.layer1[1].conv1.parameters()))
base_params = filter(lambda p: id(p) not in block1 + conv1, model.parameters())

optimizer = optim.Adam([{'params': base_params},  
                        {'params': model.layer1[0].parameters(), 'lr': 0.1}, 
                        {'params': model.layer1[1].conv1.parameters(), 'lr': 0.2}], lr=0.5)

参考

1. Access weights of a specific module in nn.Sequential()
2. pytorch学习系列（6）：不同的层设置不同的学习率（迁移学习）
3. https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-optim/