BERT生成50万句句子嵌入提速及批量处理实现问题咨询

问题解答 #

1. 加速处理的可行方法 #

• 优先使用批量推理，GPU的并行计算能力只有批量输入才能完全发挥，单句跑会浪费90%以上的算力
• 有GPU的话将模型和输入张量都移到GPU上（model.to('cuda')），推理速度比CPU提升至少10倍；没有GPU也可以开启多进程分批处理，速度也比单句循环快
• 可使用半精度推理（torch.float16/torch.bfloat16），显存占用减半的同时速度能提升30%~50%，对嵌入生成的准确率影响几乎可以忽略
• 50万条数据建议按批拆分处理，每批处理的句子数量根据显存大小调整（32/64/128都可），避免一次加载数据过多爆显存

2. 批量输入确实能大幅提升效率 #

BERT本身原生支持批量输入，单句推理时GPU大部分计算单元都处于闲置状态，在不爆显存的前提下批量越大，单位时间处理的句子数越多。3090级别的GPU每批跑128条的话，50万条句子几个小时就能处理完成，单句跑可能需要数天。 #

完整批量处理实现方案 #

你之前批量运行报错是因为没有给不同长度的句子做填充（padding）和截断，调用tokenizer时开启填充、截断参数，同时传入attention_mask让模型忽略填充部分即可，不需要单条循环处理。另外你之前求平均时没有过滤填充的token，会导致生成的嵌入不准，下面的代码也做了对应优化：

#!pip install transformers
import torch
from transformers import BertTokenizer, BertModel

# 加载分词器和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased', output_hidden_states = True)
model.eval()

# 有GPU可开启下面两行，移到GPU加速
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model = model.to(device)

# 你的全量句子可以替换到这个列表里
all_sentences = [
    "Hello I'm a single sentence",
    "And another sentence",
    "And the very very last one"
]
storage = []
# 根据显存大小调整批大小，32/64/128都可
batch_size = 32

# 按批拆分全量句子循环处理
for idx in range(0, len(all_sentences), batch_size):
    batch = all_sentences[idx:idx+batch_size]
    # 分词时自动做填充、截断，直接返回张量格式
    encoded_inputs = tokenizer(
        batch,
        padding=True,
        truncation=True,
        max_length=128, # 根据你句子的平均长度调整，太长可以适当缩小
        return_tensors='pt'
    )
    # 有GPU可开启下面一行，将输入移到GPU
    # encoded_inputs = {k: v.to(device) for k, v in encoded_inputs.items()}
    
    with torch.no_grad():
        # 整批输入直接传入模型，无需循环
        outputs = model(**encoded_inputs)
        hidden_states = outputs[2]
    
    # 取倒数第二层输出，形状为 [batch_size, 序列长度, 768]
    token_vecs = hidden_states[-2]
    # 扩展attention_mask维度，用来过滤填充的token
    attention_mask = encoded_inputs['attention_mask'].unsqueeze(-1).expand(token_vecs.size()).float()
    # 仅对非填充的token求平均得到句子嵌入，形状为 [batch_size, 768]
    sentence_embeddings = torch.sum(token_vecs * attention_mask, dim=1) / torch.clamp(attention_mask.sum(dim=1), min=1e-9)
    
    # 存储结果，不需要GPU的话可转成CPU张量/ numpy数组
    for text, embed in zip(batch, sentence_embeddings):
        storage.append((text, embed.cpu()))

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