Pytorch如何获取BERT模型最后一层隐藏状态的CLS的embedding？

遇到问题

BERT模型中最后一层的句子的CLS的embedding怎么获取？
来源于阅读 An Interpretability Illusion for BERT这篇论文

We began by creating embeddings for the 624,712 sentences
in our four datasets. To do this, we used the BERT-base
uncased model from the HuggingFace Transformers library
with no fine tuning or dataset specific modifications. We
used the final layer hidden state of each sentence’s [CLS]
token as its embedding.这最后一句话用代码实现

解决方法

经询问GPT4o，获得的答案如下：

你想要使用 BERT 模型的最后一层隐藏状态作为句子的嵌入（embedding），可以通过以下代码来实现。这段代码使用 transformers 库中的 BERT 模型来处理句子，并提取 [CLS] token 的嵌入。

from transformers import BertTokenizer, BertModel
import torch

# 加载预训练的 BERT 模型和 tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 将模型设置为评估模式
model.eval()

# 输入的句子
sentence = "This is an example sentence."

# 对输入的句子进行编码
inputs = tokenizer(sentence, return_tensors='pt')

# 获取模型的输出
with torch.no_grad():
    outputs = model(**inputs)

# 获取最后一层的隐藏状态 (outputs.last_hidden_state)
last_hidden_state = outputs.last_hidden_state

# 获取 [CLS] token 的嵌入 (最后一层隐藏状态的第一个 token)
cls_embedding = last_hidden_state[:, 0, :]

print(cls_embedding)

代码说明：

1. 加载 BERT 模型和 Tokenizer:

   • BertTokenizer 和 BertModel 分别用于对文本进行标记化和加载预训练的 BERT 模型。
   • 我们使用了 bert-base-uncased 版本的模型。

2. 设置模型为评估模式:

   • model.eval() 将模型设置为评估模式，关闭 dropout 和 batch normalization。

3. 编码输入句子:

   • tokenizer(sentence, return_tensors='pt') 将句子转换为模型所需的输入格式。return_tensors='pt' 表示输出为 PyTorch 张量。

4. 获取模型输出:

   • model(**inputs) 返回一个 BaseModelOutputWithPoolingAndCrossAttentions 对象，其中包含模型的隐藏状态。
   • outputs.last_hidden_state 是最后一层的隐藏状态，形状为 (batch_size, sequence_length, hidden_size)。

5. 提取 [CLS] token 的嵌入:

   • [CLS] token 的嵌入对应于最后一层隐藏状态的第一个 token (last_hidden_state[:, 0, :])，其形状为 (batch_size, hidden_size)。

6. 打印嵌入:

   • cls_embedding 是输入句子的嵌入。

对切片的解释

对上面一句切片的解释：

在 cls_embedding = last_hidden_state[:, 0, :] 这行代码中，[:, 0, :] 是一种高级的切片（slicing）语法，用于从 last_hidden_state 张量中提取特定的元素。让我们逐个解析这部分代码：

last_hidden_state 的维度

假设 last_hidden_state 是一个形状为 (batch_size, sequence_length, hidden_size) 的三维张量：

• batch_size: 表示一批次中句子的数量（批次大小）。
• sequence_length: 表示每个句子中的 token 数量。
• hidden_size: 表示 BERT 模型的隐藏层大小（通常是 768 维度）。

切片语法解释

• :: 代表取所有该维度的元素。
• 0: 代表取该维度的第一个元素（索引从 0 开始）。
• :: 再次代表取所有该维度的元素。

因此，[:, 0, :] 分别应用于 batch_size, sequence_length, 和 hidden_size 这三个维度。

具体解释

1. :（第一个维度）:

   • 代表取 batch_size 维度上的所有元素，即包含每个句子的结果。

2. 0（第二个维度）:

   • 代表取 sequence_length 维度上的第一个元素，也就是每个句子中的第一个 token 的隐藏状态。在 BERT 的输出中，第一个 token 通常是 [CLS] token。

3. :（第三个维度）:

   • 代表取 hidden_size 维度上的所有元素，也就是 [CLS] token 的整个隐藏状态向量。

总结

• last_hidden_state[:, 0, :] 提取了每个句子的 [CLS] token 的隐藏状态向量，其中 batch_size 维度保持不变。
• 最终得到的 cls_embedding 是一个形状为 (batch_size, hidden_size) 的二维张量，其中每一行是对应句子的 [CLS] token 的嵌入向量。

后记

2024年8月21日12点43分于上海，基于GPT4o大模型生成。