☰
BERT句向量生成实战:池化策略、微调与语义检索避坑指南
2026/10/8 7:33:03 网站建设 项目流程

简介:这份资源面向自然语言处理初学者与需要快速落地文本向量化的开发者,围绕如何调用BERT预训练模型生成句向量与词向量展开,帮助解决文本语义表示与下游任务特征提取的问题。压缩包共3个文件,均为Python脚本,整体约4KB,其中配置脚本负责指定模型权重与词汇表路径,向量提取脚本分别实现句子级与词级向量生成,另含模型结构相关代码,便于按模块理解与二次修改。资源已有1127人学习下载,说明其在入门实践中具有一定参考价值。读者可借此掌握以[CLS]标记提取整句语义、按位置获取单词上下文向量的具体做法,并将所得向量用于文本分类、情感分析、相似度计算或问答系统等场景,同时理解PyTorch或TensorFlow环境下transformers库的调用方式,为后续研究预训练模型提供可复用的脚本基础。

1. 用 BERT 预训练模型生成句向量:为什么你跑出来的向量“没那味儿”

很多团队第一次把 BERT 接进检索或聚类流程时,都会经历同一个场景:模型加载成功、forward不报错、向量也拿到了,但拿去做相似度匹配,结果和关键词匹配差不多,甚至更差。问题往往不在 BERT 本身,而在于“用哪一层输出、怎么池化、要不要微调”这三件事没定清楚。这个标题讲的就是把预训练 BERT 变成可用的句向量或词向量:词向量取隐层状态,句向量靠池化策略或 Sentence-BERT 这类结构,最终服务于语义检索、去重、聚类、召回粗排。适合已经会调 HuggingFacetransformers、但向量质量一直上不去的工程师,也适合想从 TF-IDF 迁移到语义向量的搜索团队。

2. 词向量与句向量:BERT 的哪一层输出才是你要的

2.1 词向量不是查表,而是上下文相关的隐状态

Word2Vec、GloVe 这类静态词向量,一个词只有一个向量,“苹果”在“吃苹果”和“苹果发布会”里完全一样。BERT 的输出是每个 token 在每个位置上的隐状态,同一个词在不同句子里向量不同,这就是常说的 contextual embedding。取词向量的常见做法是拿last_hidden_state,再按 tokenizer 的offset_mapping把 subword 合并回原词。注意 BERT 用的是 WordPiece,一个中文词可能被切成多个 token,直接取第一个 subword 会丢信息,常见做法是对同一词的所有 subword 取平均。

from transformers import BertTokenizer, BertModel import torch tokenizer = BertTokenizer.from_pretrained("bert-base-chinese") model = BertModel.from_pretrained("bert-base-chinese") model.eval() text = "苹果发布了新款手机" inputs = tokenizer(text, return_tensors="pt", return_offsets_mapping=True) offsets = inputs.pop("offset_mapping")[0] # 每个 token 对应原文的字符区间 with torch.no_grad(): outputs = model(**inputs) hidden = outputs.last_hidden_state[0] # [seq_len, hidden_size] # 按字符区间把 subword 合并成词向量 word_vecs = {} for idx, (start, end) in enumerate(offsets.tolist()): if start == end: # 特殊符号 [CLS] [SEP] 跳过 continue word = text[start:end] word_vecs.setdefault(word, []).append(hidden[idx]) word_vecs = {w: torch.stack(v).mean(0) for w, v in word_vecs.items()} print(word_vecs["苹果"].shape) # torch.Size([768])

逻辑说明:offset_mapping是 tokenizer 给出的字符级映射,用它把 subword 还原成原词是最稳的方式。参数上return_offsets_mapping=True必须开,否则拿不到区间;model.eval()和torch.no_grad()一起用,避免 dropout 干扰和显存浪费。bert-base-chinese的 hidden_size 是 768,如果你的下游任务对维度敏感,可以在池化后再接一层线性降维。

2.2 句向量的四种池化策略与选型

句向量不是 BERT 原生输出的东西,得自己“造”。常见四种做法:

策略做法适用场景注意点
CLS 池化取[CLS]位置向量分类微调后的模型原生 BERT 的 CLS 未针对相似度训练,效果一般
平均池化对所有 token 向量取均值通用、稳定要排除 padding,否则被 0 拉偏
最大池化逐维取最大值关键词敏感任务对噪声敏感
加权池化按 IDF 或 attention 权重加权检索场景实现复杂,收益不稳定

我一般先用平均池化做基线,再决定要不要上 Sentence-BERT。原生 BERT 的句向量在语义相似度任务上表现平平,这是有论文验证过的:未经对比学习的 BERT,句向量甚至不如 GloVe 平均。所以如果你的场景是语义检索,直接上text2vec-base-chinese或paraphrase-multilingual-MiniLM这类已经用对比学习训过的模型,比在原生 BERT 上折腾池化划算得多。

def mean_pooling(hidden, attention_mask): # hidden: [batch, seq_len, dim] mask = attention_mask.unsqueeze(-1).float() # [batch, seq_len, 1] summed = (hidden * mask).sum(dim=1) # 屏蔽 padding counts = mask.sum(dim=1).clamp(min=1e-9) return summed / counts inputs = tokenizer(["今天天气不错", "今天阳光很好"], return_tensors="pt", padding=True, truncation=True, max_length=128) with torch.no_grad(): out = model(**inputs) emb = mean_pooling(out.last_hidden_state, inputs["attention_mask"]) # 归一化后再算余弦相似度 emb = torch.nn.functional.normalize(emb, p=2, dim=1) print((emb[0] @ emb[1]).item())

参数说明:padding=True配合attention_mask是必须的,否则 batch 内短句会被 padding 污染。max_length=128对多数中文句子够用,长文档要截断或分段。归一化后余弦相似度退化成点积,省一次计算,检索时常用。

2.3 微调与否:决定向量质量的分水岭

如果你的语料是法律、医疗、工单这类垂直领域,原生 BERT 的向量空间和你的业务语义有偏差。这时候两条路:一是用 Sentence-BERT 的对比学习框架在自己的句对上微调,二是用text2vec这类已经在大规模中文语料上训过的模型直接推理。前者成本高但上限高,后者上手快。判断标准很简单:拿 100 对业务句对做人工标注,算一下原生模型的相似度排序和人工排序的 Spearman 相关系数,低于 0.5 就考虑微调。

3. 从零跑通句向量生成:环境、代码与批量推理

3.1 环境准备与模型选择

pip install transformers torch sentence-transformers # 国内下载慢可以设镜像 export HF_ENDPOINT=https://hf-mirror.com

模型选择上,中文场景我一般按这个顺序试:shibing624/text2vec-base-chinese(轻量、快)、BAAI/bge-base-zh-v1.5(检索强)、bert-base-chinese(只做基线对比)。sentence-transformers库封装了池化和归一化,直接model.encode()就能出句向量,比自己手写池化省事。

from sentence_transformers import SentenceTransformer model = SentenceTransformer("shibing624/text2vec-base-chinese") sentences = ["如何重置密码", "忘记密码怎么办", "今天天气如何"] emb = model.encode(sentences, normalize_embeddings=True, batch_size=32) print(emb.shape) # (3, 768)

normalize_embeddings=True让输出直接是单位向量,后续算余弦相似度就是点积。batch_size根据显存调,768 维模型在 8G 显存上跑 64 没问题。

3.2 批量推理与显存控制

生产环境里句子是流式来的,不能一次全加载。常见做法是分块 +DataLoader,同时用torch.cuda.amp做半精度推理。

from torch.utils.data import DataLoader def encode_in_batches(model, texts, batch_size=64, max_length=128): model.eval() all_emb = [] for i in range(0, len(texts), batch_size): batch = texts[i:i+batch_size] with torch.no_grad(), torch.cuda.amp.autocast(): emb = model.encode(batch, normalize_embeddings=True, batch_size=len(batch), show_progress_bar=False) all_emb.append(emb) import numpy as np return np.vstack(all_emb)

逻辑说明:autocast在支持 Tensor Core 的卡上能提速 30% 以上,精度损失对检索任务可忽略。np.vstack把分块结果拼回矩阵,方便后续写 FAISS 或 Milvus。注意model.encode内部已经做了 tokenize 和池化,不要再手动传attention_mask。

3.3 向量落库与相似度检索

拿到向量后,小规模(万级)直接用 numpy 算余弦,大规模上 FAISS。

import faiss import numpy as np dim = 768 index = faiss.IndexFlatIP(dim) # 内积索引,配合归一化向量等价余弦 index.add(emb.astype("float32")) query = model.encode(["密码忘了"], normalize_embeddings=True).astype("float32") scores, ids = index.search(query, k=3) print(ids, scores)

参数说明:IndexFlatIP是精确检索,数据量超过百万考虑IndexIVFFlat并调nlist。归一化向量用内积索引,省去每次算模长。如果向量没归一化,得用IndexFlatL2并做距离到相似度的转换。

4. 避坑与排查:句向量质量翻车的五个血泪现场

4.1 现象:相似句得分 0.99,不相似句也 0.95

原因:用了原生bert-base-chinese的 CLS 向量,没做对比学习,向量空间各向异性严重,所有向量挤在一个锥形区域里,余弦相似度普遍偏高。解决:换text2vec或bge这类对比学习模型,或者自己用MultipleNegativesRankingLoss微调。判断方法:随机抽 100 对无关句子算相似度,均值超过 0.7 就是各向异性问题。

4.2 现象:短句和长句相似度算出来偏低

原因:平均池化时没排除 padding,或者长句被max_length截断丢了关键信息。解决:确认attention_mask传对了,长文档按 256 或 512 分段后分别编码再聚合。我一般对长文本用“分段编码 + 平均”而不是直接截断,召回率能差 10 个点。

4.3 现象:GPU 显存爆了,batch_size 调到 1 还是 OOM

原因:max_length设太大,或者 tokenizer 没开truncation,遇到超长文本直接撑爆。解决:tokenizer(..., truncation=True, max_length=128)必须加,padding='max_length'比padding=True更耗显存,动态 padding 更省。另外检查是不是在no_grad外面跑了推理,梯度图会占大量显存。

4.4 现象:同一句话两次编码结果不一样

原因:模型没设eval(),dropout 还在生效。解决:推理前model.eval(),并用torch.no_grad()包住。如果是sentence-transformers,model.encode内部已经处理,但自己手写 forward 时容易忘。

4.5 现象:中文词向量里“的”“了”这类停用词向量模长特别大

原因:BERT 的隐状态里高频词的向量范数普遍偏大,直接算相似度会被这些词主导。解决:池化前对每个 token 向量做 L2 归一化,或者用 IDF 加权。检索场景里我一般对 query 和 doc 都做归一化,再算相似度,能压掉一部分高频词干扰。

5. 进阶技巧:用对比学习把句向量调到业务可用

原生 BERT 的句向量在垂直领域不够用,这是共识。真正让向量“有那味儿”的一步,是用业务句对做对比学习微调。核心思路是:让语义相近的句子在向量空间里靠近,不相关的推远。sentence-transformers提供了MultipleNegativesRankingLoss,训练时一个 batch 内其他样本自动当负例,不需要人工构造负样本。

from sentence_transformers import SentenceTransformer, InputExample, losses from torch.utils.data import DataLoader model = SentenceTransformer("shibing624/text2vec-base-chinese") # 每条是 (query, 正例) 句对,负例由 batch 内其他样本充当 train_examples = [ InputExample(texts=["如何重置密码", "密码忘了怎么找回"]), InputExample(texts=["订单怎么取消", "我想退掉刚下的单"]), # ... 至少几千条业务句对 ] train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32) train_loss = losses.MultipleNegativesRankingLoss(model) model.fit( train_objectives=[(train_dataloader, train_loss)], epochs=3, warmup_steps=100, output_path="./finetuned-bert-vec", optimizer_params={"lr": 2e-5}, )

参数说明:batch_size=32意味着每个 query 有 31 个负例,batch 越大负例越多,效果通常越好,但显存也涨。lr=2e-5是 BERT 微调的经典值,太大容易灾难性遗忘。epochs=3对几千条数据够用,多了会过拟合。训练完用model.encode直接出向量,和之前接口一致。

验证微调有没有效果,别只看 loss。我一般做两件事:一是拿 200 对人工标注的句对,算微调前后相似度排序和人工排序的 Spearman 相关系数;二是拿业务 query 跑一次召回,看 Top-10 里相关文档的占比。相关系数提升 0.1 以上、召回率提升 5 个点以上,这次微调就值了。如果没提升,先检查句对质量——很多团队栽在“正例其实不相似”上,标注规范比模型选择重要得多。

一个我踩过的坑:微调时用了太多“字面相似但语义不同”的句对,比如“苹果手机”和“苹果价格”,模型学出来还是字面匹配。后来改成按业务意图标注,同一意图下的不同说法才算正例,向量质量才真正上来。这个方向值不值得投入,取决于你的业务有没有稳定的句对来源;如果没有,先用现成的bge或text2vec顶着,别硬训。希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询