训练网络指定层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 match
    the 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将会被用于所 有的参数。