使用 bert-bert-chinese 预训练模型去做分类任务，这里找了新闻分类数据，数据有 20w，来自https://github.com/649453932/Bert-Chinese-Text-Classification-Pytorch/tree/master/THUCNews
数据 20w ，18w 训练数据，1w 验证数据， 1w 测试数据，10个类别我跑起来后，预测要7天7夜，于是吧数据都缩小了一些，每个类别抽一些，1800 训练数据，150 验证数据， 150 测试数据，都跑了 1.5 小时， cpu ，电脑 gpu 只有 2g 显存，带不起来
bert- base-chinses 模型下载：bert预训练模型下载-CSDN博客
训练
现在是大模型时代了，这篇文章的代码是利用大模型帮我写的的，通过大模型修正代码，并解释代码一直到可用，代码都写了注释了，整个分类流程就这样，算是一个通用模板了吧
train.py
# 导入所需的库
import torch
import os
import time
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertForSequenceClassification
from torch.optim import AdamW
# 定义数据集类，符合高内聚原则
class NewsTitleDataset(Dataset):
def __init__(self, file_path, tokenizer, max_len=128):
self.data = []
with open(file_path, 'r', encoding='utf-8') as f:
for line in f.readlines():
title, label = line.strip().split('\t')
inputs = tokenizer(title, padding='max_length', truncation=True, max_length=max_len)
self.data.append({'input_ids': inputs['input_ids'], 'attention_mask': inputs['attention_mask'], 'label': int(label)})
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
'''在使用DataLoader加载数据进行训练或验证时被调用'''
return self.data[idx]
# 训练函数（部分代码，实际训练时应包含更多细节如损失计算、模型更新等）
def train_model(model, train_loader, val_loader, optimizer, epochs=3, model_save_path='../output/bert_news_classifier'):
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 判断是否使用GPU
device = torch.device("cpu")
model.to(device)
best_val_accuracy = None # 初始化最优验证集准确率
# 创建保存目录（如果不存在）
os.makedirs(os.path.dirname(model_save_path), exist_ok=True)
# 训练几次模型
for epoch in range(epochs):
model.train() # 开启训练模式，会更新参数
for batch in train_loader:
input_ids = batch['input_ids'].to(device) # 直接通过键名访问'input_ids'
attention_mask = batch['attention_mask'].to(device) # 直接通过键名访问'attention_mask'
labels = batch['label'].to(device) # 直接通过键名访问'label'
outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss # 获取损失
optimizer.zero_grad() # 清零梯度
loss.backward() # 反向传播
optimizer.step() # 更新权重
# 在每个epoch结束时评估模型性能
model.eval()
with torch.no_grad():
val_loss = 0
correct_predictions = 0
total_samples = len(val_data) # 计算总样本数，用于后续计算准确率
for batch in val_loader:
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['label'].squeeze().to(device)
# 计算logits而不是直接获取loss
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
# logits 是模型对输入数据计算出的未归一化的类别概率分布，通常是一个形状为 (batch_size, num_classes) 的张量
logits = outputs.logits
# 手动计算loss（假设labels已转换为one-hot编码或数值标签）
# 交叉熵损失函数，多分类问题中常用的损失函数，特别适合于处理像BERT这样的预训练模型输出的logits，并且与one-hot编码的目标标签一起使用
loss_fct = torch.nn.CrossEntropyLoss()
# labels 是实际的类别标签，需要转换成一个形状为 (batch_size,) 的张量以匹配logits的展开维度。
# view(-1, model.num_labels) 会将logits展平为 (batch_size * num_classes) 的向量，使得每个样本的每个类别都有一个单独的概率值对应
loss = loss_fct(logits.view(-1, model.num_labels), labels.view(-1))
# .item() 方法用于从损失张量提取标量值。
val_loss += loss.item()
# 找出每个样本的最大概率对应的类别索引，即模型预测的结果。
# dim=1 时，表示在第二个维度上找到最大值
_, preds = torch.max(logits, dim=1)
correct_predictions += (preds == labels).sum().item()
val_accuracy = correct_predictions / total_samples
print(f'Epoch {epoch + 1}, Validation Loss: {val_loss / len(val_loader):.4f}, Accuracy: {val_accuracy * 100:.2f}%')
# 如果当前验证集上的准确率优于之前保存的最佳模型，则保存当前模型
if best_val_accuracy is None or val_accuracy > best_val_accuracy:
best_val_accuracy = val_accuracy
torch.save(model.state_dict(), model_save_path) # 保存模型参数
# 定义评估函数
def evaluate_model(model, data_loader):
device = next(model.parameters()).device
model.eval()
correct_predictions = 0
total_samples = 0
with torch.no_grad():
for batch in data_loader:
inputs = {key: batch[key].to(device) for key in ['input_ids', 'attention_mask']}
labels = batch['label'].to(device)
outputs = model(**inputs)
_, preds = torch.max(outputs.logits, dim=1)
correct_predictions += (preds == labels).sum().item()
total_samples += len(labels)
return correct_predictions / total_samples
def collate_to_tensors(batch):
input_ids = torch.tensor([example['input_ids'] for example in batch])
attention_mask = torch.tensor([example['attention_mask'] for example in batch])
labels = torch.tensor([example['label'] for example in batch])
return {'input_ids': input_ids, 'attention_mask': attention_mask, 'label': labels}
start = time.time()
# 加载预训练的tokenizer和模型
tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
with open('../data/class.txt', 'r', encoding='utf8') as f:
class_labels = f.readlines()
model = BertForSequenceClassification.from_pretrained('../bert-base-chinese', num_labels=len(class_labels)) # 假设class_labels是一个包含所有类别的列表
# 加载训练、验证和测试数据集
train_data = NewsTitleDataset('../data/train.txt', tokenizer)
val_data = NewsTitleDataset('../data/dev.txt', tokenizer)
test_data = NewsTitleDataset('../data/test.txt', tokenizer)
# 创建DataLoader，用于批处理数据
# collate_to_tensors 调用函数，保证模型接受的数据参数类型一定为 pytorch 的张量类型
# shuffle=True 于防止模型过拟合和提高泛化性能至关重要，因为它确保了模型不会因为训练数据的顺序而产生依赖性。
# batch_size 示每次迭代从数据集中取出多少个样本作为一个批次（batch）进行训练。设置合理的批量大小有助于平衡计算效率和内存使用。
train_loader = DataLoader(train_data, batch_size=32, shuffle=True, collate_fn=collate_to_tensors)
val_loader = DataLoader(val_data, batch_size=32, collate_fn=collate_to_tensors)
test_loader = DataLoader(test_data, batch_size=32, collate_fn=collate_to_tensors)
# 设置优化器与学习率
# model.parameters()：这是PyTorch中的一个方法，用于获取模型的所有可训练参数。
# lr代表学习率（Learning Rate），它是一个超参数，决定了在每个训练步骤中更新模型参数的幅度大小。给定值 2e-5 表示0.00002
optimizer = AdamW(model.parameters(), lr=2e-5)
# 开始训练
train_model(model, train_loader, val_loader, optimizer, model_save_path='../output/best_bert_news_classifier.pth')
# 测试模型（仅评估，不更新参数）
test_acc = evaluate_model(model, test_loader)
print(f'Test Accuracy: {test_acc * 100:.2f}%')
print(time.time() - start)
运行结果
预测
假如只想输入一个文本，直接得到疯了及结果，可以使用一下代码
import torch
from transformers import BertTokenizer
from transformers import BertForSequenceClassification
# 假设 model_state_dict 是从文件加载的模型参数
with open('../data/class.txt', 'r', encoding='utf8') as f:
class_labels = f.readlines()
model = BertForSequenceClassification.from_pretrained('../bert-base-chinese', num_labels=len(class_labels)) # 初始化模型结构，并指定分类类别数量
# 假设 tokenizer 是您在训练时使用的 BERT tokenizer
tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
# 加载模型参数，训练好输出的模型参数
model.load_state_dict(torch.load('../output/best_bert_news_classifier.pth'))
model.eval() # 设置模型为评估模式
def predict_news_category(model, tokenizer, text):
# 对文本进行预处理并编码
inputs = tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=128, # 根据实际情况调整最大长度
padding='max_length',
truncation=True,
return_tensors='pt'
)
input_ids = inputs['input_ids'].to(model.device)
attention_mask = inputs['attention_mask'].to(model.device)
# 将数据传递给模型以获取logits
with torch.no_grad():
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
# 获取分类结果
logits = outputs.logits
_, prediction = torch.max(logits, dim=1)
# 返回预测类别索引，实际应用中可能需要将其映射回原始类别标签
return prediction.item()
# 示例：输入一条新闻标题并预测类别
text = "车载大模型是原子弹还是茶叶蛋？"
predicted_category = predict_news_category(model, tokenizer, text)
print(f"预测的新闻类别是：{predicted_category}")
