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

Question

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

bert_output = model(input_ids, attention_mask=attention_mask)

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

Answer 1

如果您想获得所有隐藏层的输出，则需要将output_hidden_states=True kwarg添加到您的配置中。

您的代码将如下所示

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)

隐藏图层将以bert_output[2]的形式提供