别再只用BERT做语义匹配了!手把手教你用SimCSE无监督训练中文句子向量(附完整代码)
突破BERT局限:SimCSE无监督训练中文句子向量实战指南
在自然语言处理领域,语义相似度计算一直是核心任务之一。传统方法通常依赖标注数据进行监督学习,但高质量标注数据的获取成本极高。想象一下,当你接手一个客服对话分析项目,手头只有海量未标注的聊天记录,却需要快速构建一个能够准确理解用户意图的语义匹配系统——这正是SimCSE大显身手的场景。
SimCSE(Simple Contrastive Learning of Sentence Embeddings)通过创新的"Dropout即数据增强"思想,仅需无标注文本就能训练出高质量的句子向量表示。与需要复杂数据增强策略的对比学习方法不同,SimCSE的实现异常简洁,却能达到甚至超越监督学习的效果。本文将带你从零开始,用PyTorch实现一个完整的中文SimCSE训练流程。
1. 环境准备与数据加载
1.1 搭建基础环境
首先确保你的开发环境已安装必要依赖。推荐使用Python 3.8+和PyTorch 1.8+版本:
pip install torch transformers pandas tqdm scikit-learn
对于中文任务,我们选择bert-base-chinese作为基础模型。这个预训练模型在中文文本表示上已有不错的表现,适合作为SimCSE的起点。
1.2 准备无标注数据
SimCSE的美妙之处在于它对数据要求极低——只需要大量自然文本,无需任何标注。假设我们有一个客服对话数据集customer_service.txt,每行包含一条用户query或客服回复:
import pandas as pd
# 加载原始文本数据
with open('customer_service.txt', 'r', encoding='utf-8') as f:
texts = [line.strip() for line in f if len(line.strip()) > 10] # 过滤过短文本
print(f"共加载{len(texts)}条文本样例")
print("示例文本:", texts[0])
提示:数据质量对最终效果影响显著。建议预处理时去除明显噪声(如特殊符号、乱码等),但保留文本的原始表达方式,这对学习真实的语义表示很重要。
2. 模型架构设计与实现
2.1 构建SimCSE核心结构
SimCSE的核心思想非常简单:将同一个句子两次输入BERT,由于Dropout的随机性,会得到两个略有不同的向量表示,将这些表示作为正样本对,同时将其他句子的表示作为负样本。
from transformers import BertModel, BertTokenizer
import torch
import torch.nn as nn
class SimCSE(nn.Module):
def __init__(self, model_name='bert-base-chinese'):
super(SimCSE, self).__init__()
self.bert = BertModel.from_pretrained(model_name)
self.tokenizer = BertTokenizer.from_pretrained(model_name)
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.bert.to(self.device)
def forward(self, input_ids, attention_mask):
# 通过BERT获取句子表示
outputs = self.bert(input_ids, attention_mask=attention_mask)
# 使用[CLS]token的表示作为整个句子的表示
cls_rep = outputs.last_hidden_state[:, 0]
return cls_rep
2.2 实现对比损失函数
InfoNCE(Noise Contrastive Estimation)损失是SimCSE训练的关键,它鼓励正样本对的相似度高于负样本对:
def contrastive_loss(embeddings, temperature=0.05):
# embeddings: [batch_size, hidden_size]
batch_size = embeddings.size(0)
# 计算相似度矩阵
sim_matrix = torch.matmul(embeddings, embeddings.T) # [batch_size, batch_size]
# 对角线元素是正样本对的相似度
pos_sim = torch.diag(sim_matrix) / temperature
# 构造标签:对角线位置为正样本
labels = torch.arange(batch_size).to(embeddings.device)
# 计算交叉熵损失
loss = nn.CrossEntropyLoss()(sim_matrix/temperature, labels)
return loss
注意:温度参数
temperature对模型性能影响很大。较小的温度会使模型更关注困难的负样本,通常设置在0.01到0.1之间。
3. 训练流程优化技巧
3.1 高效数据加载与处理
为了充分利用GPU资源,我们需要实现一个高效的数据加载器:
from torch.utils.data import Dataset, DataLoader
class TextDataset(Dataset):
def __init__(self, texts, tokenizer, max_len=128):
self.texts = texts
self.tokenizer = tokenizer
self.max_len = max_len
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
inputs = self.tokenizer(
text,
max_length=self.max_len,
padding='max_length',
truncation=True,
return_tensors='pt'
)
return {
'input_ids': inputs['input_ids'].squeeze(),
'attention_mask': inputs['attention_mask'].squeeze()
}
# 初始化数据集和数据加载器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
dataset = TextDataset(texts, tokenizer)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
3.2 训练循环实现
下面是一个完整的训练循环,包含梯度累积和模型保存逻辑:
def train(model, dataloader, optimizer, epochs=3, accumulation_steps=4):
model.train()
total_steps = len(dataloader) * epochs
current_step = 0
for epoch in range(epochs):
for batch_idx, batch in enumerate(dataloader):
# 将数据移动到设备
input_ids = batch['input_ids'].to(model.device)
attention_mask = batch['attention_mask'].to(model.device)
# 前向传播 - 同一个batch输入两次利用Dropout得到不同表示
embeddings1 = model(input_ids, attention_mask)
embeddings2 = model(input_ids, attention_mask)
# 合并两次的结果作为最终batch
embeddings = torch.cat([embeddings1, embeddings2], dim=0)
# 计算损失
loss = contrastive_loss(embeddings)
loss = loss / accumulation_steps # 梯度累积
# 反向传播
loss.backward()
# 梯度累积到指定步数后更新参数
if (batch_idx + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
current_step += 1
if current_step % 100 == 0:
print(f"Epoch {epoch+1}, Step {current_step}/{total_steps}, Loss: {loss.item()*accumulation_steps:.4f}")
# 保存训练好的模型
torch.save(model.state_dict(), 'simcse_bert_chinese.bin')
4. 评估与应用实践
4.1 语义相似度评估
训练完成后,我们需要评估模型在语义相似度任务上的表现。一个常用的方法是计算Spearman相关系数:
from sklearn.metrics.pairwise import cosine_similarity
from scipy.stats import spearmanr
def evaluate(model, eval_pairs):
"""评估模型在语义相似度任务上的表现
eval_pairs: [(text1, text2, human_score), ...]
"""
model.eval()
pred_scores = []
human_scores = []
with torch.no_grad():
for text1, text2, score in eval_pairs:
# 编码两个句子
inputs1 = model.tokenizer(text1, return_tensors='pt', padding=True, truncation=True).to(model.device)
inputs2 = model.tokenizer(text2, return_tensors='pt', padding=True, truncation=True).to(model.device)
emb1 = model(inputs1['input_ids'], inputs1['attention_mask']).cpu().numpy()
emb2 = model(inputs2['input_ids'], inputs2['attention_mask']).cpu().numpy()
# 计算余弦相似度
sim = cosine_similarity(emb1, emb2)[0][0]
pred_scores.append(sim)
human_scores.append(score)
# 计算Spearman相关系数
spearman = spearmanr(human_scores, pred_scores).correlation
print(f"Spearman相关系数: {spearman:.4f}")
return spearman
4.2 实际应用示例
训练好的SimCSE模型可以广泛应用于各种语义相关任务。以下是一个简单的语义搜索实现:
class SemanticSearch:
def __init__(self, model, texts):
self.model = model
self.texts = texts
self.embeddings = self._encode_all(texts)
def _encode_all(self, texts):
self.model.eval()
embeddings = []
with torch.no_grad():
for text in texts:
inputs = self.model.tokenizer(text, return_tensors='pt',
padding=True, truncation=True).to(self.model.device)
emb = self.model(inputs['input_ids'], inputs['attention_mask'])
embeddings.append(emb.cpu().numpy())
return np.concatenate(embeddings, axis=0)
def search(self, query, top_k=5):
with torch.no_grad():
inputs = self.model.tokenizer(query, return_tensors='pt',
padding=True, truncation=True).to(self.model.device)
query_emb = self.model(inputs['input_ids'], inputs['attention_mask']).cpu().numpy()
# 计算余弦相似度
similarities = cosine_similarity(query_emb, self.embeddings)[0]
# 获取最相似的top_k个结果
top_indices = similarities.argsort()[-top_k:][::-1]
return [(self.texts[i], similarities[i]) for i in top_indices]
在实际项目中,我发现几个关键点对最终效果影响很大:首先是batch size要尽可能大(受限于GPU内存),因为更多的负样本能提供更好的对比信号;其次是学习率不宜过大,通常设置在1e-5到5e-5之间;最后是数据质量比数量更重要,清洗掉低质量文本能显著提升模型性能。
更多推荐


所有评论(0)