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

本文主要是介绍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)

代码说明：

加载 BERT 模型和 Tokenizer:
- BertTokenizer 和 BertModel 分别用于对文本进行标记化和加载预训练的 BERT 模型。
- 我们使用了 bert-base-uncased 版本的模型。
设置模型为评估模式:
- model.eval() 将模型设置为评估模式，关闭 dropout 和 batch normalization。
编码输入句子:
- tokenizer(sentence, return_tensors='pt') 将句子转换为模型所需的输入格式。return_tensors='pt' 表示输出为 PyTorch 张量。
获取模型输出:
- model(**inputs) 返回一个 BaseModelOutputWithPoolingAndCrossAttentions 对象，其中包含模型的隐藏状态。
- outputs.last_hidden_state 是最后一层的隐藏状态，形状为 (batch_size, sequence_length, hidden_size)。
提取 [CLS] token 的嵌入:
- [CLS] token 的嵌入对应于最后一层隐藏状态的第一个 token (last_hidden_state[:, 0, :])，其形状为 (batch_size, hidden_size)。
打印嵌入:
- 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 这三个维度。

具体解释

:（第一个维度）:
- 代表取 batch_size 维度上的所有元素，即包含每个句子的结果。
0（第二个维度）:
- 代表取 sequence_length 维度上的第一个元素，也就是每个句子中的第一个 token 的隐藏状态。在 BERT 的输出中，第一个 token 通常是 [CLS] token。
:（第三个维度）:
- 代表取 hidden_size 维度上的所有元素，也就是 [CLS] token 的整个隐藏状态向量。

总结

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