BERT的参数应该如何冻结？

BERT是极具代表性的预训练语言模型，几乎是现在很多语言理解任务，甚至是数据挖掘任务的基座。由于种种原因，部分任务在使用BERT的时候需要冻结部分参数，或者只引入部分层，因此需要对BERT中的参数进行冻结（或解冻）处理。

只引入部分层

这里的讨论全都基于12层的bert-base-uncased展开，其他模型变体也是类似的。一般来说底层transformer的输出包含更多的词语级信息，高层transformer的输出更多包含句段级的信息。如果想要只使用最下面8层transformer结构，则可以直接在初始化时进行设定。

import pytorch
from transformers import AutoConfig, AutoModel
from transformers import AutoTokenizer

bert_path = 'bert-base-uncased'
config = AutoConfig.from_pretrained(bert_path)
config.output_hidden_states = True
config.num_hidden_layers = 8

bert_layer = AutoModel.from_pretrained(bert_path, config=config)
tokenizer = AutoTokenizer.from_pretrained(bert_path, config=config)

config.num_hidden_layers = 8就是对层数的设定，通常来说都是从最下层开始往上数n层进行保留，很少有截取中间几层进行保留或者只保留最高几层的。所谓的保留，就是正常加载预训练时得到的参数，不保留的层直接将参数随机化，而且模型的运算过程也不会跟他们有关。在这样的设定下初始化BERT之后，会输出下列提示，提醒用户注意下列层并未初始化。

初始化模型时的提示

引入部分层的结果是会改变hidden_states的形状，这是因为没有初始化的层根本不参与运算，自然也就没有输出。

tys = tokenizer("this is a test line", add_special_tokens=True, truncation=True,
                                       max_length=24, padding='max_length')
bert_out = bert_layer(torch.tensor(tys["input_ids"]).unsqueeze(0), 
                      torch.tensor(tys["attention_mask"]).unsqueeze(0))
print(len(bert_out.hidden_states), bert_out.hidden_states[0].shape)
# 9 torch.Size([1, 24, 768])

需要注意的是，尽管这里这保留了8层，但是最终hidden_states中的隐藏层向量却有9个，这是因为最底层的embedding层有一个单独的输出结果也被纳入其中了。如果想取第7层的输出，代码应该写作bert_out.hidden_states[7]而不是bert_out.hidden_states[6]。

关于BERT的输出，可以参考https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertModel

冻结部分层的参数

首先来查看BERT的各层名称及是否需要梯度（冻结与否），默认所有层均需要梯度

1 2	for name ,param in bert_layer.named_parameters(): print(name, param.requires_grad)

输出如下

embeddings.word_embeddings.weight True
embeddings.position_embeddings.weight True
embeddings.token_type_embeddings.weight True
embeddings.LayerNorm.weight True
embeddings.LayerNorm.bias True
encoder.layer.0.attention.self.query.weight True
encoder.layer.0.attention.self.query.bias True
encoder.layer.0.attention.self.key.weight True
encoder.layer.0.attention.self.key.bias True
encoder.layer.0.attention.self.value.weight True
encoder.layer.0.attention.self.value.bias True
encoder.layer.0.attention.output.dense.weight True
encoder.layer.0.attention.output.dense.bias True
encoder.layer.0.attention.output.LayerNorm.weight True
encoder.layer.0.attention.output.LayerNorm.bias True
encoder.layer.0.intermediate.dense.weight True
encoder.layer.0.intermediate.dense.bias True
encoder.layer.0.output.dense.weight True
encoder.layer.0.output.dense.bias True
encoder.layer.0.output.LayerNorm.weight True
encoder.layer.0.output.LayerNorm.bias True
...
encoder.layer.11.output.LayerNorm.weight True
encoder.layer.11.output.LayerNorm.bias True
pooler.dense.weight True
pooler.dense.bias True

我们要关注的，就是所有包含encoder.layer的层，下面给出一种很简单的方式，来对特定层参数进行冻结。

1
2
3

for i in range(8):
    for p in bert_layer.encoder.layer[i].parameters():
        p.requires_grad = False

这种方法不会冻结最底层embedding层的参数，如果真的必须冻结embedding层，可以改为如下代码。

for i in range(8):
    for p in bert_layer.encoder.layer[i].parameters():
        p.requires_grad = False
for p in bert_layer.embeddings.parameters():
    p.requires_grad = False

本文作者：MyTech::Author

本文链接： https://mytech.pages.dev/2023/09/02/bert-freeze/

更新于 2024-03-26

# python # BERT # transformers