BERT词向量嵌入实战：从原理到生产环境优化

1次阅读

共计 2244 个字符，预计需要花费 6 分钟才能阅读完成。

在搜索推荐系统中，词向量嵌入的质量直接影响语义匹配的准确性。传统 Word2Vec 通过局部上下文窗口学习静态词向量，无法解决一词多义问题。而 BERT 基于 Transformer 架构，通过双向上下文建模生成动态词向量，在 ”bank” 等歧义词的场景下，能根据上下文生成不同的向量表示。

不同规模的 BERT 模型在计算资源和语义捕获能力上存在显著差异：

模型类型	参数量	显存占用	语义理解能力
BERT-base	110M	1.1GB	中等
BERT-large	340M	3.2GB	强
DistilBERT	66M	0.7GB	中等偏下

对于大多数生产场景，推荐采用 BERT-base 作为起点，在 GPU 显存不足时考虑 DistilBERT。

模型加载与异常处理

from transformers import BertModel, BertTokenizer
import torch

try:
    # 加载预训练模型和分词器
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BertModel.from_pretrained('bert-base-uncased', 
                                    output_hidden_states=True)  # 获取所有隐藏层
    model.eval()
except Exception as e:
    print(f"模型加载失败: {str(e)}")
    # 可降级到轻量级模型
    model = BertModel.from_pretrained('distilbert-base-uncased')

获取各层向量表示

def get_bert_embeddings(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    # outputs.hidden_states 包含 13 层输出（输入层 +12 个 Transformer 层）all_layers = outputs.hidden_states  
    return all_layers

向量聚合策略对比

# 策略 1：最后四层均值池化（常用方案）def mean_pooling(last_four_layers):
    # last_four_layers shape: [4, seq_len, 768]
    return torch.mean(torch.stack(last_four_layers), dim=0)

# 策略 2：CLS token 作为句子表示
def cls_pooling(last_hidden_state):
    # last_hidden_state shape: [1, seq_len, 768]
    return last_hidden_state[0][0]  # 取第一个 token（[CLS]）

FP16 量化效果实测

精度	显存占用	推理速度	语义相似度得分
FP32	1.1GB	22ms	0.872
FP16	0.6GB	15ms	0.869

启用方法：

model.half()  # 转换为 FP16

FAISS 索引构建

import faiss
import numpy as np

# 假设已有一批向量 embeddings.shape=(10000, 768)
embeddings = np.random.rand(10000, 768).astype('float32')

# 建立 IVF 索引
dimension = 768
nlist = 100  # 聚类中心数
quantizer = faiss.IndexFlatIP(dimension)
index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
index.train(embeddings)
index.add(embeddings)

# 相似度搜索
D, I = index.search(embeddings[:5], k=3)  # 查询前 5 个向量的最近 3 个邻居

必须使用中文专用 tokenizer：

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

处理长文本时建议启用自动截断：

inputs = tokenizer(text, truncation=True, max_length=512)

批量推理防 OOM 方案：

动态批次大小调整

batch_sizes = [32, 16, 8, 4]  # 降级序列
for batch_size in batch_sizes:
    try:
        process_batch(batch_size)
        break
    except RuntimeError:
        continue

梯度累积（训练时适用）

for i, batch in enumerate(batches):
    loss = model(batch).loss
    loss = loss / accumulation_steps
    loss.backward()
    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

在医疗、法律等专业领域，通用 BERT 可能无法准确捕捉术语语义。此时需要考虑：
– 领域自适应预训练（继续预训练）
– 有监督微调（如有标注数据）
– 知识增强（注入领域词典）

实际效果需要通过领域内的语义相似度任务验证，例如使用 STS- B 的领域改编版进行评估。

正文完