如何基于分块后的列表提取DistilBERT Embeddings？

针对分块后的list_train，可以通过遍历分块→逐块处理→合并结果的方式提取Embeddings，具体步骤和代码如下：

核心思路 #

将每个分块视为独立的小数据集，依次对分块内的文本做分词、模型推理，最后把所有分块的Embedding结果拼接成完整的数据集Embeddings，避免一次性处理全量数据导致内存溢出。

实现步骤与代码 #

1. 准备模型与分词器（若未初始化） #

from transformers import DistilBertTokenizer, DistilBertModel
import torch
import pandas as pd
import numpy as np

# 初始化分词器和模型
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertModel.from_pretrained('distilbert-base-uncased')
# 自动选择计算设备（GPU优先）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

2. 遍历分块提取Embeddings #

# 存储所有分块的Embeddings
all_embeddings = []

# 遍历每个分块
for chunk in list_train:
    # 从分块中提取文本列（假设文本列名为'text'，请根据实际修改）
    texts = chunk['text'].tolist()
    
    # 对当前分块的文本分词
    tokenized_inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors='pt'
    ).to(device)  # 移至对应计算设备
    
    # 模型推理，提取Embeddings（禁用梯度计算节省内存）
    with torch.no_grad():
        outputs = model(**tokenized_inputs)
        bert_embeddings = outputs.last_hidden_state
    
    # 选择句子表征方式：以下二选一即可
    # 方式1：取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token的Embedding
    cls_embeddings = bert_embeddings[:, 0, :].cpu().numpy()
    # 方式2：对所有token的Embedding做均值池化
    # cls_embeddings = torch.mean(bert_embeddings, dim=1).cpu().numpy()
    
    # 将当前分块的Embedding加入列表
    all_embeddings.append(cls_embeddings)
    
    # 清理临时变量，释放内存（GPU环境下更关键）
    del tokenized_inputs, outputs, bert_embeddings
    torch.cuda.empty_cache()

# 合并所有分块的Embeddings，得到5000条记录的完整Embedding矩阵
final_embeddings = np.concatenate(all_embeddings, axis=0)
# 可选：转为DataFrame方便后续处理
embedding_df = pd.DataFrame(final_embeddings)

关键注意事项 #

• 文本列名：代码中chunk['text']需替换为你数据集实际的文本列名称。
• 池化方式：根据任务需求选择取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token、均值池化或最大池化的Embedding作为句子表征。
• 内存优化：torch.no_grad()会禁用梯度计算，大幅降低内存占用；推理后及时删除临时变量并清理GPU缓存，可进一步避免内存溢出。

内容的提问来源于stack exchange，提问作者Prasad Joshi