模型训练-Tricks：分层学习率【①预训练模型需要设置较小学习率，②自定义的层需要设置较大学习率】

在模型的训练阶段对于不同的层我们可能会需要设置不同的学习率，尤其是当我们在预训练模型的基础上对下游任务进行微调，对于预训练模型需要设置较小的学习率，而其他一些我们自定义的层则需要较大的学习率。这篇文章通过一个简单的例子简要记录一下 Pytorch 中分层学习率的设置方式。

我们定义一个简单的 Bert 模型。

class BertTest(torch.nn.Module):
    def __init__(self,embedding_size,hidden_size):
        super(BertTest,self).__init__()
        self.bert_model = transformers.BertModel.from_pretrained('bert-base-uncased')
        self.lstm = nn.LSTM(hidden_size,embedding_size)
        self.drop_out = nn.Dropout(0.1)
        self.linear_layer = nn.Linear(embedding_size,hidden_size)

    def forward(self,inputs):
        output = self.bert(inputs)
        output = self.lstm(output)
        output = self.drop_out(output)
        output = self.linear_layer(output)
        return output

首先我们分析定义的模型结构，这个简单的模型是由12层 transformer 的 encoder 加上 lstm、linear_layer 组成。大概是这个样子

更清晰一点，我们将层数去掉，仅看模型的架构部分。

import re
def func(s):
    s = re.sub(r'\d+.','',s)
    return s
bert_test = BertTest(100,768)
params_list = [n for n,p in list(bert_test.named_parameters())]
params_list = set(map(func,params_list))
print(params_list)

好，了解完模型架构，接下来引入我们分层学习率的设计需求。

1. 预训练好的 bert 模型 lr 设定为 1e-5,其他自定义层的 lr 设定为 1e-2
2. bias、LayerNorm 层的 weight_decay 设定为 0，其他层 weight_decay 设定为 1e-2.

下面为实现代码

def get_group_parameters(model):
    params = list(model.named_parameters())
    no_decay = ['bias,','LayerNorm']
    other = ['lstm','linear_layer']
    no_main = no_decay + other

    param_group = [
        {'params':[p for n,p in params if not any(nd in n for nd in no_main)],'weight_decay':1e-2,'lr':1e-5},
        {'params':[p for n,p in params if not any(nd in n for nd in other) and any(nd in n for nd in no_decay) ],'weight_decay':0,'lr':1e-5},
        {'params':[p for n,p in params if any(nd in n for nd in other) and any(nd in n for nd in no_decay) ],'weight_decay':0,'lr':1e-2},
        {'params':[p for n,p in params if any(nd in n for nd in other) and not any(nd in n for nd in no_decay) ],'weight_decay':1e-2,'lr':1e-2},
    ]
    return param_group

optimizer = torch.optim.AdamW(param_group,lr=1e-5,eps=1e-7) # 括号中为除了分组中的默认值

我们检验一下是否所有的参数都囊括到我们的分组中：

no_decay = ['bias','LayerNorm']
other = ['lstm','linear_layer']
no_main = no_decay + other
params = list(bert_test.named_parameters())

bert_decay = ([n for n,p in params if not any(nd in n for nd in no_main)])
bert_no_decay = ([n for n,p in params if not any(nd in n for nd in other) and any(nd in n for nd in no_decay) ])
other_no_decay = ([n for n,p in params if any(nd in n for nd in other) and any(nd in n for nd in no_decay) ])
other_decay = ([n for n,p in params if any(nd in n for nd in other) and not any(nd in n for nd in no_decay) ])
all_params = [n for n,p in params]

a = set(bert_decay + bert_no_decay + other_no_decay + other_decay)
b = set(all_params)

print(a == b)

---

---

---

Pytorch 设置分层学习率 - 知乎

 【pytorch】不同层设置不同学习率_不同层用不同的学习率_阿委困的不能行的博客-CSDN博客

 分层学习率设置和学习率衰减（pytorch）_风居住的街道~的博客-CSDN博客