如何在HuggingFace Transformers库中获得经过预训练的BERT模型的中间层输出?

时间:2020-04-27 17:47:33

标签: tensorflow keras tensorflow2.0 huggingface-transformers bert-language-model

(我正在关注pytorch的BERT词嵌入教程,在该教程中,作者正在访问BERT模型的中间层。)

我想要的是使用HuggingFace的Transformers库访问TensorFlow2中BERT模型的单个输入令牌的最后4个最后层。由于每层输出的矢量长度为768,所以后4层的形状为4*768=3072(对于每个令牌)。

如何在TF / keras / TF2中实现此目的,以获得输入令牌的预训练模型的中间层? (稍后,我将尝试获取句子中每个标记的标记,但是现在一个标记就足够了。)

我正在使用HuggingFace的BERT模型:

!pip install transformers
from transformers import (TFBertModel, BertTokenizer)

bert_model = TFBertModel.from_pretrained("bert-base-uncased")  # Automatically loads the config
bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
sentence_marked = "hello"
tokenized_text = bert_tokenizer.tokenize(sentence_marked)
indexed_tokens = bert_tokenizer.convert_tokens_to_ids(tokenized_text)

print (indexed_tokens)
>> prints [7592]

输出是令牌([7592]),它应该是BERT模型的输入。

1 个答案:

答案 0 :(得分:3)

BERT模型输出的第三个元素是一个元组,它由嵌入层的输出以及中间层的隐藏状态组成。来自documentation

hidden_​​states(tuple(tf.Tensor),可选,在config.output_hidden_states=True时返回): 形状为tf.Tensor的{​​{1}}元组(一个用于嵌入的输出+一个用于每一层的输出)。

每层输出加上初始嵌入输出时模型的隐藏状态。

对于(batch_size, sequence_length, hidden_size)模型,bert-base-uncased默认为config.output_hidden_states。因此,要访问12个中间层的隐藏状态,您可以执行以下操作:

True

outputs = bert_model(input_ids, attention_mask) hidden_states = outputs[2][1:] 元组中有12个元素,从开始到最后一层都对应于所有图层,每个元素都是一个形状为hidden_states的数组。因此,例如,要访问批次中所有样本的第五个令牌的第三层的隐藏状态,您可以执行:(batch_size, sequence_length, hidden_size)


请注意,如果默认情况下加载的模型未返回隐藏状态,则可以使用hidden_states[2][:,4]类加载配置并传递BertConfig参数,如下所示:

output_hidden_state=True