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")  # 判断是否使用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}")