bert分类模型使用

使用 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")  # 判断是否使用GPUdevice = 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 = 0correct_predictions = 0total_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而不是直接获取lossoutputs = 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_samplesprint(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_accuracytorch.save(model.state_dict(), model_save_path)  # 保存模型参数# 定义评估函数
def evaluate_model(model, data_loader):device = next(model.parameters()).devicemodel.eval()correct_predictions = 0total_samples = 0with 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_samplesdef 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)# 将数据传递给模型以获取logitswith 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}")