BERT模型长文本词向量生成与内存优化实践

import torch
from transformers import AutoModel, AutoTokenizer

# 假设有一组待处理的文本数据
texts = [
    "这是一个示例文本，用于演示如何使用BERT生成词向量。",
    "BERT模型在自然语言处理任务中表现出色，但内存管理是关键。",
    # ... 更多文本
]

# 加载预训练模型和分词器
# 替换为你要使用的BERT模型名称
model_name = "bert-base-uncased" # 或者 "indolem/indobert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 如果有GPU可用，将模型移动到GPU
if torch.cuda.is_available():
    model.to('cuda')
    print("模型已加载到GPU。")
else:
    print("GPU不可用，模型将在CPU上运行。")

# 对文本进行分词，并进行填充、截断和返回PyTorch张量
tokenized_texts = tokenizer(texts,
                            max_length=512,      # 最大序列长度
                            truncation=True,     # 超过max_length的部分进行截断
                            padding=True,        # 填充到max_length
                            return_tensors='pt') # 返回PyTorch张量

print(f"分词后的输入ID形状: {tokenized_texts['input_ids'].shape}")
print(f"分词后的注意力掩码形状: {tokenized_texts['attention_mask'].shape}")

# 将输入数据移动到GPU（如果模型在GPU上）
if torch.cuda.is_available():
    input_ids = tokenized_texts['input_ids'].to('cuda')
    attention_mask = tokenized_texts['attention_mask'].to('cuda')
else:
    input_ids = tokenized_texts['input_ids']
    attention_mask = tokenized_texts['attention_mask']

# 前向传播生成词向量
with torch.no_grad():
    outputs = model(input_ids=input_ids,
                    attention_mask=attention_mask)

    # 获取最后一层隐藏状态作为词向量
    word_embeddings = outputs.last_hidden_state

# 打印词向量的形状
# 形状通常是 [batch_size, num_seq_tokens, embed_size]
# 例如：[2, 512, 768] 表示批次大小为2，序列长度为512，词向量维度为768
print(f"生成的词向量形状: {word_embeddings.shape}")

# 假设原始文本列表为 all_texts
all_texts = [
    "这是一个非常长的文本样本，可能导致内存问题...",
    # ... 2370行文本
]

# 定义一个合适的批处理大小，例如 8, 16, 32，根据GPU内存调整
batch_size = 16
all_word_embeddings = []

for i in range(0, len(all_texts), batch_size):
    batch_texts = all_texts[i:i + batch_size]

    # 对当前批次文本进行分词
    tokenized_batch = tokenizer(batch_texts,
                                max_length=512,
                                truncation=True,
                                padding=True,
                                return_tensors='pt')

    # 将输入数据移动到GPU
    if torch.cuda.is_available():
        input_ids_batch = tokenized_batch['input_ids'].to('cuda')
        attention_mask_batch = tokenized_batch['attention_mask'].to('cuda')
    else:
        input_ids_batch = tokenized_batch['input_ids']
        attention_mask_batch = tokenized_batch['attention_mask']

    # 模型前向传播
    with torch.no_grad():
        outputs_batch = model(input_ids=input_ids_batch,
                              attention_mask=attention_mask_batch)
        word_embeddings_batch = outputs_batch.last_hidden_state

    # 将当前批次的词向量添加到总列表中
    all_word_embeddings.append(word_embeddings_batch.cpu()) # 移回CPU以释放GPU内存

    # 显式清空CUDA缓存，有助于防止内存碎片化
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

# 合并所有批次的词向量
final_word_embeddings = torch.cat(all_word_embeddings, dim=0)
print(f"最终合并的词向量形状: {final_word_embeddings.shape}")