访问预训练 DistilBERT 模型多层的输出

Bul*_*aud 1 python nlp pytorch bert-language-model huggingface-transformers

我正在尝试从预训练的“DistilBERT”模型的几个不同层访问输出嵌入。(“distilbert-base-uncased”)

bert_output = model(input_ids, attention_mask=attention_mask)
Run Code Online (Sandbox Code Playgroud)

bert_output 似乎只返回输入标记最后一层的嵌入值。

eri*_*rip 5

如果你想获得所有隐藏层的输出,你需要将output_hidden_states=Truekwarg 添加到你的配置中。

你的代码看起来像

from transformers import DistilBertModel, DistilBertConfig

config = DistilBertConfig.from_pretrained('distilbert-base-cased', output_hidden_states=True)
model = DistilBertModel.from_pretrained('distilbert-base-cased', config=config)
Run Code Online (Sandbox Code Playgroud)

隐藏层将作为bert_output[2]