訪問幾層預訓練的 DistilBERT 模型的輸出

Question

我正在嘗試從預訓練的“DistilBERT”模型的幾個不同層訪問輸出嵌入。 （“蒸餾酒基未裝箱”）

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]