如何避免RAM崩溃，批量获取BERT的[CLS]句嵌入？

1. 批量处理句子，减少模型调用次数 #

单条循环处理会导致内存碎片化严重，批量处理能大幅降低内存开销。根据Colab免费版的内存情况，选择合适的batch size（比如8、16或32）：

from transformers import BertTokenizer, BertModel
import torch

# 初始化模型和分词器，切换到评估模式
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model.eval()

batch_size = 16
total_rows = len(rows_list)

for i in range(0, total_rows, batch_size):
    # 取出当前批次的句子
    batch = rows_list[i:i+batch_size]
    texts = [row['Text'] for row in batch]
    
    # 批量分词
    tokenized = tokenizer(texts, padding='max_length', max_length=512, truncation=True, return_tensors="pt")
    
    # 关闭梯度计算，节省内存
    with torch.no_grad():
        # 不需要output_hidden_states=True，避免存储所有层的状态
        outputs = model(**tokenized)
        # 提取[CLS] token的嵌入，转成numpy数组减少内存占用
        cls_embeddings = outputs.last_hidden_state[:, 0, :].numpy()
    
    # 把嵌入赋值回原字典
    for idx, row in enumerate(batch):
        row['Embedding'] = cls_embeddings[idx]
    
    # 清理当前批次的变量，释放内存
    del tokenized, outputs, cls_embeddings
    torch.cuda.empty_cache()  # GPU环境下清理显存

2. 关闭不必要的输出，砍掉冗余内存占用 #

原代码里output_hidden_states=True会保存所有隐藏层的状态，这会吃掉大量内存。如果只需要最后一层的[CLS]嵌入，完全不需要开启这个参数，直接用outputs.last_hidden_state就能拿到结果。

3. 换用更轻量的模型 #

如果内存还是吃紧，换成参数更少的BERT变体，比如DistilBERT——它的参数只有原版BERT的60%，速度快30%，内存占用低很多，语义嵌入效果也接近原版：

from transformers import DistilBertTokenizer, DistilBertModel
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertModel.from_pretrained('distilbert-base-uncased')
model.eval()

4. 及时清理内存，避免泄漏 #

• 必须用torch.no_grad()上下文管理器，禁用梯度计算（评估模式不需要梯度，能省大量内存）。
• 每个批次处理完后，手动删除不再使用的张量变量，GPU环境下调用torch.cuda.empty_cache()清理显存。
• 把PyTorch张量转成numpy数组再存入字典，numpy数组在CPU内存中占用的资源比GPU张量少得多。

5. 检查Colab运行时配置 #

确保已经切换到GPU运行时（路径：Runtime > Change runtime type > Hardware accelerator选GPU），GPU的显存比CPU内存更适合处理大模型。如果还是内存溢出，重启运行时（Runtime > Restart runtime）清理残留内存。

内容的提问来源于stack exchange，提问作者Aadithya Seshadri