高效生成BERT词嵌入：解决内存溢出挑战

import torch
from transformers import AutoModel, AutoTokenizer

# 示例输入文本列表
texts = ['这是一个测试句子，它可能有点长，但我们希望它能被正确处理。', 
         '另一个示例文本，用于演示如何生成词嵌入。']

# 加载匹配的模型和分词器
# 替换为您的模型名称，例如 "bert-base-uncased"
model_name = "indolem/indobert-base-uncased" 
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 将模型移动到GPU（如果可用）
if torch.cuda.is_available():
    model.to('cuda')
    print("模型已移至GPU。")
else:
    print("未检测到GPU，模型将在CPU上运行。")

# 对批量句子进行分词，截断至512，并进行填充
tokenized_texts = tokenizer(texts, 
                            max_length=512,       # 最大序列长度
                            truncation=True,      # 启用截断，超出max_length的部分将被截断
                            padding=True,         # 启用填充，短于max_length的部分将被填充
                            return_tensors='pt')  # 返回PyTorch张量

# 将分词结果移动到GPU（如果模型在GPU上）
if torch.cuda.is_available():
    tokenized_texts = {k: v.to('cuda') for k, v in tokenized_texts.items()}

print(f"分词结果的input_ids形状: {tokenized_texts['input_ids'].shape}")

# 前向传播
with torch.no_grad():
    input_ids = tokenized_texts['input_ids']
    attention_mask = tokenized_texts['attention_mask']

    outputs = model(input_ids=input_ids, 
                    attention_mask=attention_mask)

    # 获取最后一层的隐藏状态作为词嵌入
    word_embeddings = outputs.last_hidden_state

# 打印词嵌入的形状
print(f"生成的词嵌入形状: {word_embeddings.shape}")
# 预期输出形状示例: torch.Size([batch_size, num_seq_tokens, embed_size])
# 例如: torch.Size([2, 512, 768])

from torch.utils.data import DataLoader, TensorDataset

# 假设您有一个非常大的文本列表
all_texts = ['长文本1', '长文本2', ..., '长文本N'] # N可能非常大

# 定义批次大小
batch_size = 16 # 根据您的GPU内存调整，尝试16, 8, 4等更小的值

# 分词所有文本 (注意：如果all_texts非常大，这一步本身可能耗内存，可以考虑分批次分词)
# 为了演示方便，我们假设分词结果可以一次性存储
tokenized_inputs = tokenizer(all_texts, 
                             max_length=512, 
                             truncation=True, 
                             padding='max_length', # 确保所有批次长度一致
                             return_tensors='pt')

input_ids_tensor = tokenized_inputs['input_ids']
attention_mask_tensor = tokenized_inputs['attention_mask']

# 创建一个TensorDataset
dataset = TensorDataset(input_ids_tensor, attention_mask_tensor)

# 创建DataLoader
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

all_embeddings = []

# 迭代处理每个批次
print(f"\n开始分批处理，批次大小为: {batch_size}")
with torch.no_grad():
    for batch_idx, batch in enumerate(dataloader):
        batch_input_ids, batch_attention_mask = batch

        # 将批次数据移动到GPU
        if torch.cuda.is_available():
            batch_input_ids = batch_input_ids.to('cuda')
            batch_attention_mask = batch_attention_mask.to('cuda')

        # 模型前向传播
        outputs = model(input_ids=batch_input_ids, 
                        attention_mask=batch_attention_mask)

        # 获取词嵌入并移回CPU（可选，但推荐，以释放GPU内存）
        batch_word_embeddings = outputs.last_hidden_state.cpu()
        all_embeddings.append(batch_word_embeddings)
        print(f"  处理批次 {batch_idx+1}/{len(dataloader)}，词嵌入形状: {batch_word_embeddings.shape}")

# 合并所有批次的词嵌入
final_embeddings = torch.cat(all_embeddings, dim=0)
print(f"\n所有文本的最终词嵌入形状: {final_embeddings.shape}")