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

Yag*_*gel 5 keras tensorflow tensorflow2.0 bert-language-model huggingface-transformers

(我正在关注这个关于 BERT 词嵌入的 pytorch 教程,在教程中,作者访问了 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]
Run Code Online (Sandbox Code Playgroud)

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

tod*_*day 7

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

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

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

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

outputs = bert_model(input_ids, attention_mask)
hidden_states = outputs[2][1:]
Run Code Online (Sandbox Code Playgroud)

hidden_statestuple中有12个元素,对应从头到尾的所有层,每个元素都是一个shape数组(batch_size, sequence_length, hidden_size)。因此,例如,进入第三层的隐藏状态,所有该批次样品的第五代币,你可以这样做:hidden_states[2][:,4]


请注意,如果您正在加载的模型默认情况下不返回隐藏状态,那么您可以使用BertConfigclass 和 passoutput_hidden_state=True参数加载配置,如下所示:

config = BertConfig.from_pretrained("name_or_path_of_model",
                                    output_hidden_states=True)

bert_model = TFBertModel.from_pretrained("name_or_path_of_model",
                                         config=config)
Run Code Online (Sandbox Code Playgroud)