如何冻结T5变压器模型的部分

How to freeze parts of T5 transformer model

我知道T5每一层都有K、Q、V向量。它还有一个前馈网络。我想冻结 K、Q 和 V 向量,只在 T5 的每一层上训练前馈层。我使用 Pytorch 库。该模型可以是 huggingface T5 模型或其修改版本的包装器。我知道如何使用以下代码冻结所有参数:

tokenizer = AutoTokenizer.from_pretrained(underlying_model_name)
model = T5ForConditionalGeneration.from_pretrained(underlying_model_name)

for p in model.parameters():
    p.requires_grad = False # freezing

你能指导我怎么做吗?

这个 github project 可能会有帮助,但它适用于 Roberta 和 GPT,我可以将它改编为 T5 吗?

我已经根据 Huggingface 论坛的 this discussion 改编了一个解决方案。 基本上,您必须指定要冻结的 modules/pytorch 图层的名称。

在您的 T5 特定案例中,我首先查看了模型摘要:

from transformers import T5ModelForConditionalGeneration

model = T5ModelForConditionalGeneration.from_pretrained("t5-small")
print(model)

这给出了以下(缩写输出):

T5ForConditionalGeneration(
  (shared): Embedding(32128, 512)
  (encoder): T5Stack(
    (embed_tokens): Embedding(32128, 512)
    (block): ModuleList(
      (0): T5Block(
        (layer): ModuleList(
          (0): T5LayerSelfAttention(
            (SelfAttention): T5Attention(
              (q): Linear(in_features=512, out_features=512, bias=False)
              (k): Linear(in_features=512, out_features=512, bias=False)
              (v): Linear(in_features=512, out_features=512, bias=False)
              (o): Linear(in_features=512, out_features=512, bias=False)
              (relative_attention_bias): Embedding(32, 8)
            )
            (layer_norm): T5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (1): T5LayerFF(
            (DenseReluDense): T5DenseReluDense(
              (wi): Linear(in_features=512, out_features=2048, bias=False)
              (wo): Linear(in_features=2048, out_features=512, bias=False)
              (dropout): Dropout(p=0.1, inplace=False)
            )
            (layer_norm): T5LayerNorm()
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
      )
[...]  # abbreviated output

有了这个,我们就可以生成一个我们想要冻结的模块列表。特别是,我决定冻结编码器的整个 T5LayerSelfAttention 块(此外,解码器的 T5LayerCrossAttention 块):

# All modules in the 
modules_to_freeze = [model.encoder.block[i].layer[0] for i in range(len(model.encoder.block))]
# And the decoder modules, which has both a SelfAttention (layer[0]) 
modules_to_freeze.extend([model.decoder.block[i].layer[0] for i in range(len(model.decoder.block))])
# and CrossAttention (layer[1]) block
modules_to_freeze.extend([model.decoder.block[i].layer[1] for i in range(len(model.decoder.block))])

然后干脆冻结各个模块中的所有参数:

for module in modules_to_freeze:
    for param in module.parameters():
        param.requires_grad = False  # Actual freezing operation

您可以通过 运行 以下内容验证这些在您的模型中是否实际冻结:

for param in model.parameters():
    print(param.requires_grad)

它也应该打印出很多 False。如果你真的只想冻结 K、Q 和 V,你可以将上述过程调整为只 sub-select 你想要的模块。