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

小灰灰 2023-10-14 09:32 24阅读 0赞

在模型的训练阶段对于不同的层我们可能会需要设置不同的学习率，尤其是当我们在预训练模型的基础上对下游任务进行微调，对于预训练模型需要设置较小的学习率，而其他一些我们自定义的层则需要较大的学习率。这篇文章通过一个简单的例子简要记录一下 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 组成。大概是这个样子

![format_png][]

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

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)

![format_png 1][]

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

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)

![format_png 2][]

![format_png 3][]

--------------------

[Pytorch 设置分层学习率 - 知乎][Pytorch _ -]

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

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

[format_png]: https://img-blog.csdnimg.cn/img_convert/1a1d955bddb9e686a23cabbef4c273a8.webp?x-oss-process=image/format,png
[format_png 1]: https://img-blog.csdnimg.cn/img_convert/8b73327ca0702db8774236ace027948d.webp?x-oss-process=image/format,png
[format_png 2]: https://img-blog.csdnimg.cn/img_convert/b4051dd1d9a9010ec50957f31a644bc9.webp?x-oss-process=image/format,png
[format_png 3]: https://img-blog.csdnimg.cn/img_convert/b5cb307b0787415fbb9abdd787217295.webp?x-oss-process=image/format,png
[Pytorch _ -]: https://zhuanlan.zhihu.com/p/436849099
[pytorch_-CSDN]: https://blog.csdn.net/rocking_struggling/article/details/126253368
[pytorch_-CSDN 1]: https://blog.csdn.net/Ming_LQ/article/details/113763667