训练网络指定层pytorch实现方法

最近在研究Mask R-CNN，该网络一部分是跟Faster R-CNN（https://arxiv.org/pdf/1506.01497v3.pdf）相似的，同样的，在模型训练实现时，其中一种方法叫做交替训练（Alternating training），想利用该方法就涉及到如何对网络进行指定层的训练，今天就总结一下pytorch中的实现方法，既然写了指定层训练，那就把参数的单独设置也介绍一下，pytorch提供的这些方法提高了网络训练的灵活度，是很好用的方法。

指定层训练

每个变量都有一个标记：requires_grad允许从梯度计算中细分排除子图，并可以提高效率。当我们把参数属性设置为requires_grad=False时，该参数固定不变，不参与训练，不更新，具体操作如下：

# 载入预训练模型参数后...
for name, value in model.named_parameters():if name 满足某些条件:value.requires_grad = False# setup optimizer
params = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.Adam(params, lr=1e-4)

检查是否固定：被固定的输出False,未被固定的输出True

for name, value in model.named_parameters():print(name,value.requires_grad)  #打印所有参数requires_grad属性，True或False

将满足条件的参数的 requires_grad 属性设置为False，同时 filter 函数将模型中属性 requires_grad = True 的参数筛选出来，传到优化器（以Adam为例）中，只有这些参数会被求导数和更新。

当我们指定的参数比较多时，可以利用正则表达式来定义参与训练的层，具体实现如下：

def set_trainable(net, layer_regex, model=None, indent=0, verbose=1):"""Sets model layers as trainable if their names matchthe given regular expression."""for layer_name,param in net.named_parameters():trainable = bool(re.fullmatch(layer_regex, layer_name))  # re.fullmatch()返回一个和模式串完全匹配的字符串if not trainable:param.requires_grad = False#将与layer_regex不匹配的层固定layer_regex = {# all layers but the backbone"heads": r"(fpn.P5\_.*)|(fpn.P4\_.*)|(fpn.P3\_.*)|(fpn.P2\_.*)|(rpn.*)|(classifier.*)|(mask.*)",# From a specific Resnet stage and up"3+": r"(fpn.C3.*)|(fpn.C4.*)|(fpn.C5.*)|(fpn.P5\_.*)|(fpn.P4\_.*)|(fpn.P3\_.*)|(fpn.P2\_.*)|(rpn.*)|(classifier.*)|(mask.*)","4+": r"(fpn.C4.*)|(fpn.C5.*)|(fpn.P5\_.*)|(fpn.P4\_.*)|(fpn.P3\_.*)|(fpn.P2\_.*)|(rpn.*)|(classifier.*)|(mask.*)","5+": r"(fpn.C5.*)|(fpn.P5\_.*)|(fpn.P4\_.*)|(fpn.P3\_.*)|(fpn.P2\_.*)|(rpn.*)|(classifier.*)|(mask.*)",# All layers"all": ".*",
}#用正则表达式标识要训练的层if layers in layer_regex.keys():layers = layer_regex[layers]self.set_trainable(layers)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)model.train_model(dataset_train, dataset_val,learning_rate=config.LEARNING_RATE,epochs=40,layers='heads')

通过选择layers来选择我们要训练更新的层。

为每个参数单独设置选项

当我们想指定每一层的学习率时，可以这样操作：

optim.SGD([{'params': model.base.parameters()},{'params': model.classifier.parameters(), 'lr': 1e-3}], lr=1e-2, momentum=0.9)

这意味着model.base的参数将会使用1e-2的学习率，model.classifier的参数将会使用1e-3的学习率，并且0.9的momentum将会被用于所有的参数。