如何提取BERT隐藏状态用于Logistic Regression等传统分类算法

1. 关于Softmax层的疑问澄清 #

你混淆了两个完全不同的Softmax层：

• 文献中提到的要移除的是分类头最末尾用于输出分类概率的Softmax层，和Transformer注意力模块内部用于计算注意力权重的Softmax没有任何关联
• 你完全不需要获取Softmax处理前的attention_scores，这个是注意力计算的中间变量，不是文献中提到的BERT隐藏状态特征

2. 提取BERT隐藏状态作为特征的实现方法 #

你只需要在加载模型时开启隐藏状态输出，推理时提取对应层的输出即可，具体步骤和代码示例如下：

2.1 环境依赖准备 #

需要提前安装好transformers、torch、scikit-learn、numpy四个依赖库。

2.2 核心代码实现 #

from transformers import CamembertTokenizer, CamembertForSequenceClassification
import torch
import numpy as np
from sklearn.linear_model import LogisticRegression

# 1. 加载预训练模型和分词器，开启隐藏状态输出
tokenizer = CamembertTokenizer.from_pretrained("你的预训练模型路径/开源模型名")
model = CamembertForSequenceClassification.from_pretrained(
    "你的预训练模型路径/开源模型名",
    output_hidden_states=True, # 核心参数，开启隐藏状态输出开关
    ignore_mismatched_sizes=True # 如果是你自己微调过的分类模型可以加这个参数避免维度报错
)
model.eval() # 切换到评估模式，关闭dropout等训练专属机制

# 2. 批量提取文本特征的函数
def extract_bert_features(text_list, max_len=128, batch_size=16):
    all_features = []
    # 分batch处理避免显存溢出
    for idx in range(0, len(text_list), batch_size):
        batch_text = text_list[idx:idx+batch_size]
        # 文本编码
        inputs = tokenizer(
            batch_text,
            padding=True,
            truncation=True,
            max_length=max_len,
            return_tensors="pt"
        )
        # 关闭梯度计算，节省显存提升推理速度
        with torch.no_grad():
            outputs = model(**inputs)
            # hidden_states结构：(模型层数, batch大小, 序列长度, 隐藏层维度)
            # 取最后一层的<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>标记输出作为整句的特征向量
            last_layer_hidden = outputs.hidden_states[-1]
            cls_feature = last_layer_hidden[:, 0, :].cpu().numpy()
            all_features.append(cls_feature)
    # 把所有batch的特征拼接成完整矩阵
    return np.concatenate(all_features, axis=0)

# 3. 特征提取+逻辑回归训练预测示例
# 假设你有自己的训练文本、训练标签、测试文本
train_features = extract_bert_features(train_text_list)
test_features = extract_bert_features(test_text_list)

# 训练逻辑回归分类器
lr_clf = LogisticRegression(max_iter=1000)
lr_clf.fit(train_features, train_label_list)

# 测试集预测
pred_labels = lr_clf.predict(test_features)

3. 效果优化建议 #

• 如果单一层特征效果不好，可以尝试取倒数2-4层的<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>特征做拼接或者加权平均，通常能获得更好的特征表达
• 如果是长文本任务，可以适当调大max_len参数，保留更多文本信息

内容的提问来源于stack exchange，提问作者AI-bobobo