Bert项目--新闻标题文本分类

技术细节

1、下载模型

2、config文件

[3、BERT 文本分类数据预处理流程](#3、BERT 文本分类数据预处理流程)

技术细节

1、下载模型

文件名称--a0_download_model.py

使用 ModelScope 库从模型仓库下载 BERT-base-Chinese 预训练模型，并将其保存到本地指定目录。

python 复制代码

# 模型下载
from modelscope import snapshot_download

model_dir = snapshot_download('google-bert/bert-base-chinese', local_dir=r"D:\src\bert-base-chinese")

2、config文件

数据加载与保存的路径

python 复制代码

# 根目录
        self.root_path = 'E:/PythonLearning/full_mask_project2/'
        # 原始数据路径
        self.train_datapath = self.root_path + '01-data/train.txt'
        self.test_datapath = self.root_path + '01-data/test.txt'
        self.dev_datapath = self.root_path + '01-data/dev.txt'
        # 类别文档
        self.class_path = self.root_path + "01-data/class.txt"
        # 类别名列表
        self.class_list = [line.strip() for line in open(self.class_path, encoding="utf-8")]  # 类别名单

        # 模型训练保存路径
        self.model_save_path = self.root_path + "dm_03_bert/save_models/bert_classifer_model.pt"  # 模型训练结果保存路径

加载预训练Bert模型以及其分词器和配置文件

python 复制代码

        self.bert_path = r"E:\PythonLearning\full_mark_Project\dm04_Bert\bert-base-chinese"  # 预训练BERT模型的路径
        self.bert_model = BertModel.from_pretrained(self.bert_path)  # 加载预训练BERT模型
        self.tokenizer = BertTokenizer.from_pretrained(self.bert_path)  # BERT模型的分词器
        self.bert_config = BertConfig.from_pretrained(self.bert_path)  # BERT模型的配置

设置模型参数

python 复制代码

        # todo 参数
        self.num_classes = len(self.class_list)  # 类别数
        self.num_epochs = 1  # epoch数
        self.batch_size = 64  # mini-batch大小
        self.pad_size = 32  # 每句话处理成的长度(短填长切)
        self.learning_rate = 5e-5  # 学习率

完整config代码

python 复制代码

import torch
import datetime
from transformers import BertModel, BertTokenizer, BertConfig

# 获取当前日期字符串
current_date = datetime.datetime.now().date().strftime("%Y%m%d")


# 配置类
class Config(object):
    def __init__(self):
        """
        配置类，包含模型和训练所需的各种参数。
        """
        self.model_name = "bert"  # 模型名称
        # todo 路径
        # 根目录
        self.root_path = 'E:/PythonLearning/full_mask_project2/'
        # 原始数据路径
        self.train_datapath = self.root_path + '01-data/train.txt'
        self.test_datapath = self.root_path + '01-data/test.txt'
        self.dev_datapath = self.root_path + '01-data/dev.txt'
        # 类别文档
        self.class_path = self.root_path + "01-data/class.txt"
        # 类别名列表
        self.class_list = [line.strip() for line in open(self.class_path, encoding="utf-8")]  # 类别名单

        # 模型训练保存路径
        self.model_save_path = self.root_path + "dm_03_bert/save_models/bert_classifer_model.pt"  # 模型训练结果保存路径

        # 模型训练+预测的时候  训练设备，如果GPU可用，则为cuda，否则为cpu
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.bert_path = r"E:\PythonLearning\full_mark_Project\dm04_Bert\bert-base-chinese"  # 预训练BERT模型的路径
        self.bert_model = BertModel.from_pretrained(self.bert_path)  # 加载预训练BERT模型
        self.tokenizer = BertTokenizer.from_pretrained(self.bert_path)  # BERT模型的分词器
        self.bert_config = BertConfig.from_pretrained(self.bert_path)  # BERT模型的配置

        # todo 参数
        self.num_classes = len(self.class_list)  # 类别数
        self.num_epochs = 1  # epoch数
        self.batch_size = 64  # mini-batch大小
        self.pad_size = 32  # 每句话处理成的长度(短填长切)
        self.learning_rate = 5e-5  # 学习率

        # TODO 量化模型存放地址
        # 注意: 量化的时候模型需要的设备首选是cpu
        self.bert_model_quantization_model_path = self.root_path + "dm_03_bert/save_models/bert_classifer_quantization_model.pt"  # 模型训练结果保存路径


if __name__ == '__main__':
    # 测试
    conf = Config()
    print(conf.device)
    print(conf.class_list)
    print(conf.tokenizer)
    input_size = conf.tokenizer.convert_tokens_to_ids(["你", "好", "中", "人"])
    print(input_size)
    print(conf.bert_model)
    print(conf.bert_config)

3、BERT 文本分类数据预处理流程

数据加载

python 复制代码

def load_raw_data(file_path):
    """
    从指定文件中加载原始数据。处理文本文件，返回(文本, 标签类别索引)元组列表
    参数:file_path: 原始文本文件路径
    返回:list: 包含(文本, 标签类别索引)的元组列表，类别为int类型
    [('体验2D巅峰 倚天屠龙记十大创新概览', 8), ('60年铁树开花形状似玉米芯(组图)', 5)]
    """
    result = []
    # 打印指定文件
    with open(file_path, 'r', encoding='utf-8') as f:
        # 使用tqdm包装文件读取迭代器，以便显示加载数据的进度条
        for line in tqdm(f, desc=f"加载原始数据{file_path}"):
            # 移除行两端的空白字符
            line = line.strip()
            # 跳过空行
            if not line:
                continue
            # 将行分割成文本和标签两部分
            text, label = line.split("\t")
            # 将标签转为int作为类别
            label = int(label)
            # 将文本和转换为整数的标签作为元组添加到数据列表中
            result.append((text, label))
    # 返回处理后的列表
    return result

数据集构建

python 复制代码

# todo 2.自定义数据集
class TextDataset(Dataset):
    # 初始化数据
    def __init__(self, data_list):
        self.data_list = data_list

    # 返回数据集长度
    def __len__(self):
        return len(self.data_list)

    # 根据样本索引,返回对应的特征和标签
    def __getitem__(self, idx):
        text, label = self.data_list[idx]
        return text, label

批处理（padding/collate）

复制代码

每当 DataLoader 从 Dataset 中取出一批batch 的原始数据后，
就会调用 collate_fn 来对这个 batch 进行统一处理（如填充、转换为张量等）。

batch 是一个包含多个 (文本, 标签) 元组的列表

使用 zip(*batch) 将元组列表"转置"为两个元组：一个包含所有文本，一个包含所有标签

例如：[("text1", 1), ("text2", 2)] → ("text1", "text2") 和 (1, 2)

add_special_tokens=True: 自动添加 [CLS] 和 [SEP] 等特殊token

padding='max_length': 将所有文本填充到固定长度 conf.pad_size

max_length=conf.pad_size: 设置最大长度

truncation=True: 如果文本超过最大长度则截断

return_attention_mask=True: 返回注意力掩码

input_ids: 文本转换为的数字token ID序列

attention_mask: 标记哪些位置是实际文本(1)，哪些是填充部分(0)

python 复制代码

def collate_fn(batch):
    """
    对batch数据进行padding处理
    参数: batch: 包含(文本, 标签)元组的batch数据
    返回: tuple: 包含处理后的input_ids, attention_mask和labels的元组
    """
    # todo 使用zip()将一批batch数据中的(text, label)元组拆分成两个独立的元组
    # texts = [item[0] for item in batch]
    # labels = [item[1] for item in batch]
    texts, labels = zip(*batch)
    # 对文本进行padding
    text_tokens = conf.tokenizer.batch_encode_plus(
        texts,
        add_special_tokens=True,  # 默认True,自动添加 [CLS] 和 [SEP]
        # padding=True,自动填充到数据中的最大长度       padding='max_length':填充到指定的固定长度
        padding='max_length',
        max_length=conf.pad_size,  # 设定目标长度
        truncation=True,  # 开启截断，防止超出模型限制
        return_attention_mask=True  # 请求返回注意力掩码，以区分输入中的有效信息和填充信息
    )
    # 从文本令牌中提取输入ID
    input_ids = text_tokens['input_ids']
    # 从文本令牌中提取注意力掩码
    attention_mask = text_tokens['attention_mask']
    # 将输入的token ID列表转换为张量
    input_ids = torch.tensor(input_ids)
    # 将注意力掩码列表转换为张量
    attention_mask = torch.tensor(attention_mask)
    # 将标签列表转换为张量
    labels = torch.tensor(labels)
    # 返回转换后的张量元组
    return input_ids, attention_mask, labels

DataLoader 封装

python 复制代码

def build_dataloader():
    # 加载原始数据
    train_data_list = load_raw_data(conf.train_datapath)
    dev_data_list = load_raw_data(conf.dev_datapath)
    test_data_list = load_raw_data(conf.test_datapath)

    # 构建训练集
    train_dataset = TextDataset(train_data_list)
    dev_dataset = TextDataset(dev_data_list)
    test_dataset = TextDataset(test_data_list)

    # 构建DataLoader
    train_dataloader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=False, collate_fn=collate_fn)
    dev_dataloader = DataLoader(dev_dataset, batch_size=conf.batch_size, shuffle=False, collate_fn=collate_fn)
    test_dataloader = DataLoader(test_dataset, batch_size=conf.batch_size, shuffle=False, collate_fn=collate_fn)

    return train_dataloader, dev_dataloader, test_dataloader

完整代码

python 复制代码

# 加载数据工具类
from tqdm import tqdm
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, TextDataset
from transformers.utils import PaddingStrategy
from config import Config
import time

# 加载配置
conf = Config()


# todo 加载并处理原始数据
def load_raw_data(file_path):
    """
    从指定文件中加载原始数据。处理文本文件，返回(文本, 标签, 类别)元组列表
    参数:file_path: 文本文件路径
    返回:list: 包含(文本, 标签, 类别)的元组列表，类别为int类型
    """
    result = []
    # 打印指定文件
    with open(file_path, 'r', encoding='utf-8') as f:
        # 使用tqdm包装文件读取迭代器，以便显示加载数据的进度条
        for line in tqdm(f, desc="加载数据..."):
            # 移除行两端的空白字符
            line = line.strip()
            # 跳过空行
            if not line:
                continue
            # 将行分割成文本和标签两部分
            text, label = line.split("\t")
            # 将标签转为int作为类别
            label = int(label)
            # 将文本和转换为整数的标签作为元组添加到数据列表中
            result.append((text, label))
    # 返回处理后的列表
    return result


# todo 自定义数据集
class TextDataset(Dataset):
    # 初始化数据
    def __init__(self, data_list):
        self.data_list = data_list

    # 返回数据集长度
    def __len__(self):
        return len(self.data_list)

    # 根据样本索引,返回对应的特征和标签
    def __getitem__(self, idx):
        text, label = self.data_list[idx]
        return text, label


# todo 批量处理数据
# 每当 DataLoader 从 Dataset 中取出一个 batch 的原始数据后，
# 就会调用 collate_fn 来对这个 batch 进行统一处理（如填充、转换为张量等）。
def collate_fn(batch):
    """
    对batch数据进行padding处理
    参数: batch: 包含(文本, 标签)元组的batch数据
    返回: tuple: 包含处理后的input_ids, attention_mask和labels的元组
    """
    # todo 使用zip()将一批batch数据中的(text, label)元组拆分成两个独立的元组
    # texts = [item[0] for item in batch]
    # labels = [item[1] for item in batch]
    texts, labels = zip(*batch)
    # 对文本进行padding
    text_tokens = conf.tokenizer.batch_encode_plus(
        texts,
        add_special_tokens=True,  # 默认True,自动添加 [CLS] 和 [SEP]
        padding='max_length',  # 固定长度
        max_length=conf.pad_size,  # 设定目标长度
        truncation=True,  # 开启截断，防止超出模型限制
        return_attention_mask=True  # 请求返回注意力掩码，以区分输入中的有效信息和填充信息
    )
    # 从文本令牌中提取输入ID
    input_ids = text_tokens['input_ids']
    # 从文本令牌中提取注意力掩码
    attention_mask = text_tokens['attention_mask']
    # 将输入的token ID列表转换为张量
    input_ids = torch.tensor(input_ids)
    # 将注意力掩码列表转换为张量
    attention_mask = torch.tensor(attention_mask)
    # 将标签列表转换为张量
    labels = torch.tensor(labels)
    # 返回转换后的张量元组
    return input_ids, attention_mask, labels


# todo 构建dataloader
def build_dataloader():
    # 加载原始数据
    train_data_list = load_raw_data(conf.train_datapath)
    dev_data_list = load_raw_data(conf.dev_datapath)
    test_data_list = load_raw_data(conf.test_datapath)

    # 构建训练集
    train_dataset = TextDataset(train_data_list)
    dev_dataset = TextDataset(dev_data_list)
    test_dataset = TextDataset(test_data_list)

    # 构建DataLoader
    train_dataloader = DataLoader(train_dataset, batch_size=conf.batch_size, shuffle=False, collate_fn=collate_fn)
    dev_dataloader = DataLoader(dev_dataset, batch_size=conf.batch_size, shuffle=False, collate_fn=collate_fn)
    test_dataloader = DataLoader(test_dataset, batch_size=conf.batch_size, shuffle=False, collate_fn=collate_fn)

    return train_dataloader, dev_dataloader, test_dataloader


if __name__ == '__main__':
    # 测试load_raw_data方法
    data_list = load_raw_data(conf.dev_datapath)
    print(data_list[:10])

    # 测试TextDataset类
    dataset = TextDataset(data_list)
    print(dataset[0])
    print(dataset[1])

    # 测试build_dataloader方法
    train_dataloader, dev_dataloader, test_dataloader = build_dataloader()
    print(len(train_dataloader))
    print(len(dev_dataloader))
    print(len(test_dataloader))

    # 测试collate_fn方法
    """
    for i, batch in enumerate(train_dataloader)流程如下:
    1.DataLoader 从你的 Dataset 中取出一组索引；
    2.使用这些索引调用 Dataset.__getitem__ 获取原始样本；
    3.将这一组样本组成一个 batch（通常是 (text, label) 元组的列表）；
    4.自动调用你传入的 collate_fn 函数来处理这个 batch 数据；
    5.返回处理后的 batch（如 input_ids, attention_mask, labels）供模型使用。
    """
    for i, batch in enumerate(train_dataloader):
        print(len(batch))
        print(i)
        input_ids, attention_mask, labels = batch
        # print("input_ids: ", input_ids.tolist())
        print("input_ids.shape: ", input_ids.shape)
        # print("attention_mask: ", attention_mask.tolist())
        print("attention_mask.shape: ", attention_mask.shape)
        # print("labels: ", labels.tolist())
        print("labels.shape: ", labels.shape)
        break

4、对输入文本进行分类

流程：
1. 加载预训练 BERT 模型（BertModel）作为特征提取器。
2. 添加全连接层（nn.Linear）进行类别预测。
3. 提供测试代码，演示从文本输入到预测结果的全过程。

outputs 是一个包含多个输出的元组，对于分类任务我们主要关注：

last_hidden_state: 序列中每个 token 的隐藏状态 (shape: [batch_size, sequence_length, hidden_size])

pooler_output: 经过池化的整个序列的表示 (shape: [batch_size, hidden_size])

假设：

batch_size = 32

sequence_length = 128 (经过 padding/truncation 后的长度)

hidden_size = 768 (BERT-base 的隐藏层大小)

num_classes = 10

数据流过程：

输入 input_ids: [32, 128]

输入 attention_mask: [32, 128]

BERT 输出 pooler_output: [32, 768]

全连接层输出 logits: [32, 10]

python 复制代码

import torch
import torch.nn as nn
from transformers import BertModel
from config import Config

# 加载配置
config = Config()


# 定义bert模型
class BertClassifier(nn.Module):
    def __init__(self):
        # 初始化父类类的构造函数
        super().__init__()
        # 下面的BertModel是从transformers库中加载的预训练模型
        # config.bert_path是预训练模型的路径
        self.bert = BertModel.from_pretrained(config.bert_path)
        # 定义全连接层（fc），用于分类任务
        # 输入尺寸是Bert模型隐藏层的大小，即768（对于Base模型）
        # 输出尺寸是类别数量，由config.num_classes指定
        self.fc = nn.Linear(config.bert_config.hidden_size, config.num_classes)

    def forward(self, input_ids, attention_mask):
        # 使用BERT模型处理输入的token ID和注意力掩码，获取BERT模型的输出
        # outputs是: _,pooled
        outputs = self.bert(
            input_ids=input_ids,  # 输入的token ID
            attention_mask=attention_mask  # 注意力掩码用于区分有效token和填充token
        )
        # print(outputs) # 观察结果
        # 通过全连接层对BERT模型的输出进行分类
        logits = self.fc(outputs.pooler_output)
        # 返回分类的logits（未归一化的预测分数）
        return logits


# 测试以上模型
if __name__ == '__main__':
    # 测试
    model = BertClassifier()
    # 加载
    from transformers import BertTokenizer

    tokenizer = BertTokenizer.from_pretrained(config.bert_path)

    # 示例文本
    texts = ["我喜欢你", "今天天气真好"]

    # 编码文本
    encoded_inputs = tokenizer(
        texts,
        # padding=True,  #  所有的填充到文本最大长度
        padding="max_length", # 所有的填充到指定的max_length长度
        truncation=True, # 如果超出指定的max_length长度，则截断
        max_length=10,
        return_tensors="pt"  # 返回 pytorch 张量,"pt" 时，分词器会将输入文本转换为模型可接受的格式
    )

    # 获取 input_ids 和 attention_mask
    input_ids = encoded_inputs["input_ids"]
    attention_mask = encoded_inputs["attention_mask"]
    print('input_ids:', input_ids)
    print('attention_mask:', attention_mask)
    print('======================================')
    # 预测
    logits = model(input_ids, attention_mask)
    print(logits)  # 每一行对应一个样本，每个数字表示该样本属于某一类别的"得分"（logit），没有经过 softmax 归一化。
    print('-------------------------------')
    # 获取预测概率
    probs = torch.softmax(logits, dim=-1)
    print(probs)  # 归一化后该样本属于某类的概率（范围在 0~1 之间）,概率最高的就是预测结果
    print('-------------------------------')
    preds = torch.argmax(probs, dim=-1)
    print(preds)  # 得到每个样本的预测类别。表示两个输入文本被模型预测为类别 6（从 0 开始计数）。

5、计算模型的分类性能指标

功能：传入模型、数据、设备，返回分类报告

完整代码：

python 复制代码

import torch
from sklearn.metrics import classification_report, f1_score, accuracy_score, precision_score, recall_score
from tqdm import tqdm


def model2dev(model, data_loader, device):
    """
    在验证或测试集上评估 BERT 分类模型的性能。
    参数：
        model (nn.Module): BERT 分类模型。
        data_loader (DataLoader): 数据加载器（验证或测试集）。
        device (str): 设备（"cuda" 或 "cpu"）。
    返回：
        tuple: (分类报告, F1 分数, 准确度, 精确度，召回率)
            - report: 分类报告（包含每个类别的精确度、召回率、F1 分数等）。
            - f1score: 微平均 F1 分数。
            - accuracy: 准确度。
            - precision: 微平均精确度
            - recall: 微平均召回率
    """
    # todo 1. 设置模型为评估模式（禁用 dropout,并改变batch_norm行为）
    model.eval()
    # 2. 初始化列表，存储预测结果和真实标签
    all_preds, all_labels = [], []
    # 3. todo torch.no_grad()禁用梯度计算以提高效率并减少内存占用
    with torch.no_grad():
        # 4. 遍历数据加载器，逐批次进行预测
        for i, batch in enumerate(tqdm(data_loader, desc="验证集评估中...")):
            # 4.1 提取批次数据并移动到设备
            input_ids, attention_mask, labels = batch
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)
            labels = labels.to(device)
            # 4.2 前向传播：模型预测
            outputs = model(input_ids, attention_mask=attention_mask)

            # 4.3 获取预测结果（最大 logits分数 对应的类别）
            y_pred_list = torch.argmax(outputs, dim=1)

            # 4.4 存储预测和真实标签
            all_preds.extend(y_pred_list.cpu().tolist())
            all_labels.extend(labels.cpu().tolist())

    # 5. 计算分类报告、F1 分数、准确率，精确率，召回率
    report = classification_report(all_labels, all_preds)
    f1score = f1_score(all_labels, all_preds, average='macro')
    accuracy = accuracy_score(all_labels, all_preds)
    precision = precision_score(all_labels, all_preds, average='macro')
    recall = recall_score(all_labels, all_preds, average='macro')

    # 6. 返回评估结果
    return report, f1score, accuracy, precision, recall

6、模型训练

口诀：15241

加载配置文件和参数（1）

python 复制代码

# todo 加载配置对象，包含模型参数、路径等
conf = Config()
# todo 导入数据处理工具类
from a1_dataloader_utils import build_dataloader
# todo 导入bert模型
from a2_bert_classifer_model import BertClassifier

准备数据（4）

python 复制代码

    # todo 1、准备数据
    train_dataloader, dev_dataloader, test_dataloader = build_dataloader()

    # todo 2、准备模型
    # 2.1初始化bert分类模型
    model = BertClassifier()
    # 2.2将模型移动到指定的设备
    model.to(conf.device)

    # todo 3.准备损失函数
    loss_fn = nn.CrossEntropyLoss()

    # todo 4.准备优化器
    optimizer = AdamW(model.parameters(), lr=conf.learning_rate)

模型训练

外层遍历轮次，内层遍历批次

前向传播：模型预测、计算损失

反向传播：梯度清零、计算梯度、参数更新

python 复制代码

# todo 5.开始训练模型
    # 初始化F1分数，用于保存最好的模型
    best_f1 = 0.0
    # todo 5.1 外层循环遍历每个训练轮次
    #  （每次需要设置训练模式，累计损失，预存训练集测试和真实标签）
    for epoch in range(conf.num_epochs):
        # 设置模型为训练模式
        model.train()
        # 初始化累计损失，初始化训练集预测和真实标签
        total_loss = 0.0
        train_preds, train_labels = [], []
        # todo 5.2 内层循环遍历训练DataLoader每个批次
        for i, batch in enumerate(tqdm(train_dataloader, desc="训练集训练中...")):
            # 提取批次数据并移动到设备
            input_ids, attention_mask, labels = batch
            input_ids = input_ids.to(conf.device)
            attention_mask = attention_mask.to(conf.device)
            labels = labels.to(conf.device)

            # todo 前向传播：模型预测
            logits = model(input_ids, attention_mask)
            # todo 计算损失
            loss = loss_fn(logits, labels)
            # 累计损失
            total_loss += loss.item()
            # todo 获取预测结果（最大logits对应的类别）
            y_pred_list = torch.argmax(logits, dim=1)
            # todo 存储预测和真实标签，用于计算训练集指标
            train_preds.extend(y_pred_list.cpu().tolist())
            train_labels.extend(labels.cpu().tolist())

            # todo 梯度清零
            optimizer.zero_grad()
            # todo 反向传播：计算梯度
            loss.backward()
            # todo 参数更新：根据梯度更新模型参数
            optimizer.step()

验证评估

python 复制代码

            # todo 每10个批次或一个轮次结束，计算训练集指标
            if (i + 1) % 10 == 0 or i == len(train_dataloader) - 1:
                # 计算准确率和f1值
                acc = accuracy_score(train_labels, train_preds)
                f1 = f1_score(train_labels, train_preds, average='macro')
                # 获取batch_count，并计算平均损失
                batch_count = i % 10 + 1
                avg_loss = total_loss / batch_count
                # todo 打印训练信息
                print(f"\n轮次: {epoch + 1}, 批次: {i + 1}, 损失: {avg_loss:.4f}, acc准确率:{acc:.4f}, f1分数:{f1:.4f}")
                # todo 清空累计损失和预测和真实标签
                total_loss = 0.0
                train_preds, train_labels = [], []

            # todo 每100个批次或一个轮次结束，计算验证集指标，打印，保存模型
            if (i + 1) % 100 == 0 or i == len(train_dataloader) - 1:
                # 计算在测试集的评估报告，准确率，精确率，召回率，f1值
                report, f1score, accuracy, precision, recall = model2dev(model, dev_dataloader, conf.device)
                print("验证集评估报告：\n", report)
                print(f"验证集的f1: {f1score:.4f}, accuracy:{accuracy:.4f}, precision:{precision:.4f}, recall:{recall:.4f}")
                # todo 将模型再设置为训练模式
                model.train()
                # todo 如果验证F1分数优于历史最佳，保存模型
                if f1score > best_f1:
                    # 更新历史最佳F1分数
                    best_f1 = f1score
                    # 保存模型
                    torch.save(model.state_dict(), conf.model_save_path)
                    print("保存模型成功, 当前f1分数:", best_f1)

完整代码

python 复制代码

import torch
import torch.nn as nn
from torch.optim import AdamW
from sklearn.metrics import f1_score, accuracy_score
from tqdm import tqdm
from model2dev_utils import model2dev
from config import Config
from a1_dataloader_utils import build_dataloader
# 忽略的警告信息
import warnings

warnings.filterwarnings("ignore")
# todo 导入bert模型
from a2_bert_classifer_model import BertClassifier

# 加载配置对象，包含模型参数、路径等
conf = Config()



def model2train():
    """
    训练 BERT 分类模型并在验证集上评估，保存最佳模型。
    参数：无显式参数，所有配置通过全局 conf 对象获取。
    返回：无返回值，训练过程中保存最佳模型到指定路径。
    """
    # todo 1、准备数据
    train_dataloader, dev_dataloader, test_dataloader = build_dataloader()

    # todo 2、准备模型
    # 2.1初始化bert分类模型
    model = BertClassifier()
    # 2.2将模型移动到指定的设备
    model.to(conf.device)

    # todo 3.准备损失函数
    loss_fn = nn.CrossEntropyLoss()

    # todo 4.准备优化器
    optimizer = AdamW(model.parameters(), lr=conf.learning_rate)

    # todo 5.开始训练模型
    # 初始化F1分数，用于保存最好的模型
    best_f1 = 0.0
    # todo 5.1 外层循环遍历每个训练轮次
    #  （每次需要设置训练模式，累计损失，预存训练集测试和真实标签）
    for epoch in range(conf.num_epochs):
        # 设置模型为训练模式
        model.train()
        # 初始化累计损失，初始化训练集预测和真实标签
        total_loss = 0.0
        train_preds, train_labels = [], []
        # todo 5.2 内层循环遍历训练DataLoader每个批次
        for i, batch in enumerate(tqdm(train_dataloader, desc="训练集训练中...")):
            # 提取批次数据并移动到设备
            input_ids, attention_mask, labels = batch
            input_ids = input_ids.to(conf.device)
            attention_mask = attention_mask.to(conf.device)
            labels = labels.to(conf.device)

            # todo 前向传播：模型预测
            outputs = model(input_ids, attention_mask)
            # todo 计算损失
            loss = loss_fn(outputs, labels)
            # todo 累计损失
            total_loss += loss.item()
            # todo 获取预测结果（最大logits对应的类别）
            y_pred_list = torch.argmax(outputs, dim=1)
            # todo 存储预测和真实标签，用于计算训练集指标
            train_preds.extend(y_pred_list.cpu().tolist())
            train_labels.extend(labels.cpu().tolist())

            # todo 梯度清零
            optimizer.zero_grad()
            # todo 反向传播：计算梯度
            loss.backward()
            # todo 参数更新：根据梯度更新模型参数
            optimizer.step()

            # todo 每10个批次或一个轮次结束，计算训练集指标
            if (i + 1) % 10 == 0 or i == len(train_dataloader) - 1:
                # 计算准确率和f1值
                acc = accuracy_score(train_labels, train_preds)
                f1 = f1_score(train_labels, train_preds, average='macro')
                # 获取batch_count，并计算平均损失
                batch_count = i % 10 + 1
                avg_loss = total_loss / batch_count
                # todo 打印训练信息
                print(f"\n轮次: {epoch + 1}, 批次: {i + 1}, 损失: {avg_loss:.4f}, acc准确率:{acc:.4f}, f1分数:{f1:.4f}")
                # todo 清空累计损失和预测和真实标签
                total_loss = 0.0
                train_preds, train_labels = [], []

            # todo 每100个批次或一个轮次结束，计算验证集指标，打印，保存模型
            if (i + 1) % 100 == 0 or i == len(train_dataloader) - 1:
                # 计算在测试集的评估报告，准确率，精确率，召回率，f1值
                report, f1score, accuracy, precision, recall = model2dev(model, dev_dataloader, conf.device)
                print("验证集评估报告：\n", report)
                print(f"验证集f1: {f1score:.4f}, accuracy:{accuracy:.4f}, precision:{precision:.4f}, recall:{recall:.4f}")
                # 将模型设置为训练模式
                model.train()
                # todo 如果验证F1分数优于历史最佳，保存模型
                if f1score > best_f1:
                    # 更新历史最佳F1分数
                    best_f1 = f1score
                    # 保存模型
                    torch.save(model.state_dict(), conf.model_save_path)
                    print("保存模型成功, 当前f1分数:", best_f1)





if __name__ == '__main__':
    model2train()

7、基于BERT的文本分类预测接口

加载预训练好的BERT分类模型（BertClassifier）。
对输入文本进行分词、编码、转换为张量。
使用模型预测文本类别。
返回类别名称（如 "education"）

python 复制代码

import torch
from a2_bert_classifer_model import BertClassifier
from config import Config

# 加载配置
conf = Config()

# todo 准备模型
model = BertClassifier()
# todo 加载模型参数
model.load_state_dict(torch.load(conf.model_save_path))
# todo 添加模型到指定设备
model.to(conf.device)
# todo 设置模型为评估模式
model.eval()


# TODO 定义predict_fun函数预测函数
def predict_fun(data_dict):
    """
    根据用户录入数据,返回分类信息
    :param 参数 data_dict: {"text":"状元心经：考前一周重点是回顾和整理"}
    :return: 返回 data_dict: {"text":"状元心经：考前一周重点是回顾和整理", "pred_class":"education"}
    """
    # todo 获取文本
    text = data_dict['text']
    # todo 将文本转为id
    text_tokens = conf.tokenizer.batch_encode_plus(
        [text],
        padding="max_length",
        max_length=conf.pad_size,
        pad_to_max_length=True
    )
    # todo 获取input_ids和attention_mask
    input_ids = text_tokens['input_ids']
    attention_mask = text_tokens['attention_mask']
    # todo 将input_ids和attention_mask转为tensor, 并指定到设备
    input_ids = torch.tensor(input_ids).to(conf.device)
    attention_mask = torch.tensor(attention_mask).to(conf.device)

    # todo 设置不进行梯度计算(在该上下文中禁用梯度计算，提升推理速度并减少内存占用)
    with torch.no_grad():
        # 前向传播(模型预测)
        output = model(input_ids, attention_mask)
        # 获取预测类别索引张量
        output = torch.argmax(output, dim=1)
        # 获取预测类别索引标量
        pred_idx = output.item()
        # 获取类别名称
        pred_class = conf.class_list[pred_idx]
        print(pred_class)
        # 将预测结果添加到data_dict中
        data_dict['pred_class'] = pred_class
    # 返回data_dict
    return data_dict


if __name__ == '__main__':
    data_dict = {'text': '状元心经：考前一周重点是回顾和整理'}
    print(predict_fun(data_dict))

问题总结

持续更新中......