BERT 简单的文本情感二分类
约 941 字大约 3 分钟
2025-12-20
1 导入库并设置环境
相关信息
导入必要的Python库：
检测并设置计算设备：如果有GPU（CUDA），则使用GPU以加速计算；否则使用CPU。
定义模型：指定BERT模型为"google-bert/bert-base-chinese"，这是一个预训练的中文BERT模型，用于序列分类任务。
import torch
from datasets import load_dataset
from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
from transformers import DataCollatorWithPadding
import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
# 检测是否有 GPU，如果有则使用，否则使用 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 模型名称
MODEL_NAME = "google-bert/bert-base-chinese"2 加载数据集
相关信息
从Hugging Face平台加载名为lansinuote/ChnSentiCorp的中文情感分析数据集。
该数据集已预分为以下子集：
训练集（train）：用于训练模型。
验证集（validation）：用于调整超参数和监控训练过程。
测试集（test）：用于最终评估模型性能。
# 加载数据集
dataset = load_dataset("lansinuote/ChnSentiCorp")
print("正在运行 Full Training")
# 使用完整数据集
train_dataset = dataset['train']
val_dataset = dataset['validation']
test_dataset = dataset['test']
print(f"数据量 -> 训练: {len(train_dataset)}, 验证: {len(val_dataset)}, 测试: {len(test_dataset)}")3 数据预处理
相关信息
加载"google-bert/bert-base-chinese"预训练模型对应的分词器。
定义预处理函数，将文本转换为模型可接受的数字ID形式，并自动截断至128长度以提高处理效率。
对训练集、验证集和测试集分别进行批量分词处理。
创建数据整理器，在训练时动态补零，确保同一批次中的文本长度一致并节省内存。
tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
def preprocess_function(examples):
# max_length=128 对于评论通常够用了，太长会增加显存负担
return tokenizer(examples["text"], truncation=True, max_length=128)
# 对数据集进行批量处理
tokenized_train = train_dataset.map(preprocess_function, batched=True)
tokenized_val = val_dataset.map(preprocess_function, batched=True)
tokenized_test = test_dataset.map(preprocess_function, batched=True)
# 数据整理器（动态 Padding，节省显存）
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)4 定义评估指标
相关信息
接收训练过程中生成的预测结果和真实标签。 从预测逻辑中选择概率最高的类别作为最终预测。
- 准确率（accuracy）：正确预测的比例
- F1分数（f1）：精确率和召回率的调和平均
- 精确率（precision）：预测为正例中实际为正例的比例
- 召回率（recall）：实际正例中被正确预测的比例
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')
acc = accuracy_score(labels, predictions)
return {
'accuracy': acc,
'f1': f1,
'precision': precision,
'recall': recall
}5 初始化模型与训练参数
相关信息
接收训练过程中生成的预测结果和真实标签。 从预测逻辑中选择概率最高的类别作为最终预测。
- 准确率（accuracy）：正确预测的比例
- F1分数（f1）：精确率和召回率的调和平均
- 精确率（precision）：预测为正例中实际为正例的比例
- 召回率（recall）：实际正例中被正确预测的比例
# 加载模型，num_labels=2 代表二分类
model = BertForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)
model.to(device)
training_args = TrainingArguments(
output_dir="./full_model_output",
learning_rate=2e-5, # 经典的 BERT 微调学习率
per_device_train_batch_size=16, # 显存如果只有 4G/8G，这里可能要改小成 8
per_device_eval_batch_size=32,
num_train_epochs=3, # 标准训练轮数
weight_decay=0.01,
evaluation_strategy="epoch", # 每个 epoch 验证一次
save_strategy="epoch", # 每个 epoch 保存一次
load_best_model_at_end=True, # 只要最好的那个模型
metric_for_best_model="accuracy",
logging_steps=100, # 日志不用太频繁
push_to_hub=False,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_train,
eval_dataset=tokenized_val,
tokenizer=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics,
)6 开始训练
# 开始训练
print("开始训练")
trainer.train()7 模型推理
print("\n------ 最终测试集评估 (Full) ------")
results = trainer.predict(tokenized_test)
for key, value in results.metrics.items():
print(f"{key}: {value:.4f}")8 模型保存
trainer.save_model("./final_sentiment_bert")
tokenizer.save_pretrained("./final_sentiment_bert")
print("模型已保存至 ./final_sentiment_bert")