PyTorch中HuggingFace SentenceTransformer微调无更新问题排查

核心问题分析 #

• 优化器绑定错误：初始化optimizer时用了model.parameters()，但实际要微调的是embedding_model，导致优化器根本没更新目标模型的参数。
• encode方法默认关闭梯度：SentenceTransformer的encode方法默认切换到eval模式，且自动禁用梯度计算（torch.no_grad()），生成的embeddings没有梯度信息，反向传播无法更新模型参数。
• 手动创建tensor破坏梯度链：把cosine_similarities重新包装成torch.tensor(..., requires_grad=True)，切断了和embedding_model的梯度关联，即使计算loss，梯度也传不到原模型。
• 学习率设置不合理：AdamW用0.1的学习率过大，会导致训练不稳定甚至参数不更新。

修复后的完整代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化模型并切换到训练模式
embedding_model = SentenceTransformer('all-MiniLM-L12-v2').to(device)
embedding_model.train()

# 绑定正确的模型参数，设置合理学习率
optimizer = torch.optim.AdamW(embedding_model.parameters(), lr=2e-5)
loss_function = nn.BCEWithLogitsLoss()

# 训练循环
for epoch in range(10):
    total_loss = 0.0
    embedding_model.train()  # 每个epoch确保处于训练模式
    for batch in train_dataloader:
        text_data1, text_data2, labels = batch['text_data1'], batch['text_data2'], batch['labels']
        labels = labels.to(device).float()

        # 直接调用模型forward方法获取embeddings，保留梯度链
        embeddings1 = embedding_model(text_data1)
        embeddings2 = embedding_model(text_data2)
        
        # 计算余弦相似度，适配损失输入格式
        cosine_similarities = F.cosine_similarity(embeddings1, embeddings2, dim=-1)
        # 可选：添加线性层将[-1,1]的相似度映射为无界logits，提升损失适配性
        # logits = nn.Linear(1, 1).to(device)(cosine_similarities.unsqueeze(1)).squeeze()
        logits = cosine_similarities

        # 计算损失并反向传播
        loss = loss_function(logits, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    
    print(f"Epoch {epoch+1}, Average Loss: {total_loss/len(train_dataloader):.4f}")

额外优化建议 #

• 添加线性层适配损失：余弦相似度范围是[-1,1]，而BCEWithLogitsLoss期望无界的logits输入，建议在余弦相似度后加nn.Linear(1,1)层，将输出映射到合适范围，训练效果会更稳定。
• 梯度裁剪：若出现梯度爆炸，可添加torch.nn.utils.clip_grad_norm_(embedding_model.parameters(), max_norm=1.0)稳定训练。
• 学习率调度：使用torch.optim.lr_scheduler.StepLR或CosineAnnealingLR动态调整学习率，提升收敛效果。
• 验证集监控：加入验证集，每个epoch后评估性能，避免过拟合。

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