📚 课程目录第2章微调BERT进行情感分类

2. 微调BERT进行情感分类

任务说明

本案例使用Hugging Face Transformers库加载预训练的bert-base-uncased模型，在IMDb电影评论数据集上进行微调，实现正面/负面情感二分类，并在测试集上评估准确率。核心目标是跑通完整的微调闭环：数据加载→分词→训练→评估。

环境依赖

需要安装以下Python包（建议Python 3.9+）：

pip install transformers datasets torch scikit-learn

完整实现代码

以下代码完成数据加载、分词、模型微调与评估全流程。使用IMDb数据集的子集（训练2000条、测试500条）以保证在普通GPU上数分钟内完成训练；如需完整数据，移除num_samples参数即可。

import torch
from torch.utils.data import DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AdamW,
    get_scheduler,
)
from datasets import load_dataset
from sklearn.metrics import accuracy_score
from tqdm.auto import tqdm

# 1. 加载数据集（使用子集加速演示）
dataset = load_dataset("imdb", split="train").shuffle(seed=42).select(range(2000))
test_dataset = load_dataset("imdb", split="test").shuffle(seed=42).select(range(500))

# 2. 加载分词器与模型
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 3. 分词处理
def tokenize_fn(batch):
    return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=256)

train_enc = dataset.map(tokenize_fn, batched=True)
test_enc = test_dataset.map(tokenize_fn, batched=True)

train_enc.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
test_enc.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

train_loader = DataLoader(train_enc, batch_size=16, shuffle=True)
test_loader = DataLoader(test_enc, batch_size=32)

# 4. 训练配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
optimizer = AdamW(model.parameters(), lr=2e-5)
num_epochs = 3
num_training_steps = num_epochs * len(train_loader)
lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps,
)

# 5. 微调训练
model.train()
progress_bar = tqdm(range(num_training_steps))
for epoch in range(num_epochs):
    for batch in train_loader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)
        progress_bar.set_postfix({"loss": loss.item()})

# 6. 评估
model.eval()
all_preds = []
all_labels = []
for batch in test_loader:
    batch = {k: v.to(device) for k, v in batch.items()}
    with torch.no_grad():
        outputs = model(**batch)
    logits = outputs.logits
    preds = torch.argmax(logits, dim=-1)
    all_preds.extend(preds.cpu().tolist())
    all_labels.extend(batch["labels"].cpu().tolist())

acc = accuracy_score(all_labels, all_preds)
print(f"Test Accuracy: {acc:.4f}")

运行方式

将上述代码保存为finetune_bert.py，在终端执行：

python finetune_bert.py

首次运行会自动下载约400MB的BERT模型权重与IMDb数据集，请确保网络畅通。训练过程在消费级GPU（如RTX 3060）上约需3～5分钟，纯CPU环境约需20～30分钟。

验证方法与通过标准

代码执行完毕后，终端会打印Test Accuracy数值。判定标准：

程序无报错完整跑通，且打印出准确率数值；
准确率应明显高于随机猜测的50%。使用2000条训练数据微调3轮，预期准确率在85%～92%区间（该数值为基于模型能力的合理预期，非实测结果）；
若使用完整IMDb训练集（25000条）并训练3轮，准确率通常可达92%以上。

如需验证模型确实学到了情感语义，可自行构造两条测试文本（如"This movie was fantastic!"与"A waste of time."），用训练好的模型预测其标签（1为正面，0为负面）。

常见问题

显存不足（OOM）：将batch_size从16调低至8或4，同时将max_length从256调低至128。
数据集下载失败：IMDb数据集托管在Hugging Face Hub，若网络受限可设置镜像export HF_ENDPOINT=https://hf-mirror.com后重试。
训练不收敛：确认标签映射正确（IMDb中label 1=正面，0=负面），且学习率使用2e-5～5e-5范围内的BERT微调推荐值。