bert新闻标题分类

使用 bert 完成文本分类任务，数据有 20w，来自https://github.com/649453932/Bert-Chinese-Text-Classification-Pytorch/tree/master/THUCNews

下载即可：

模型使用 bert-base-chinese 下载参考：bert预训练模型下载-CSDN博客

实现了新闻分类，小编在这做个笔记，整个流程也就是对 bert 模型的应用，写了注释，方便学习查看，把代码放这里记录一下：

import os
import torch
from transformers import (
    get_linear_schedule_with_warmup,BertTokenizer,
    AdamW,
    AutoModelForSequenceClassification,
    AutoConfig
)
from torch.utils.data import DataLoader, dataset
import time
import numpy as np
from sklearn import metrics
from datetime import timedelta


data_dir = 'THUCNews/data'
# 在代码开始部分添加全局变量和函数
global_batch_size = 4  # 初始化为固定的batch_size
max_batch_size = 32  # 根据实际情况设置最大允许的batch_size


def get_optimal_batch_size():
    global global_batch_size
    # 这里是检查GPU可用内存并尝试增大批次大小的逻辑
    # 具体实现可能需要根据您的设备和任务进行调整
    # 以下仅为模拟示例，实际操作时请替换为正确的方法
    free_memory = torch.cuda.memory_allocated() / (1024**3)  # 获取当前GPU空闲显存（单位：GB）
    optimal_bs = min(max_batch_size, int(free_memory * 0.8))  # 按80%的空闲显存分配批次大小
    if optimal_bs > global_batch_size:
        global_batch_size = optimal_bs
    return global_batch_size


def read_file(path):
    with open(path, 'r', encoding="UTF-8") as file:
        docus = file.readlines()
        newDocus = []
        for data in docus:
            newDocus.append(data)
    return newDocus


class Label_Dataset(dataset.Dataset):  # 建立自定义数据集
    def __init__(self, data):
        self.data = data

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

    def __getitem__(self, ind):
        onetext = self.data[ind]
        content, label = onetext.split('\t')
        label = torch.LongTensor([int(label)])
        return content, label


# 读取数据内容
trainContent = read_file(os.path.join(data_dir, "train.txt"))
testContent = read_file(os.path.join(data_dir, "test.txt"))
# 封成数据类型
traindataset = Label_Dataset(trainContent)
testdataset = Label_Dataset(testContent)
# 封装成数据加载器
testdataloder = DataLoader(testdataset, batch_size=1, shuffle=False)
batch_size = 1
traindataloder = DataLoader(traindataset, batch_size=get_optimal_batch_size(), shuffle=True)
# 加载器类别名称
class_list = [x.strip() for x in open(
    os.path.join(data_dir, "class.txt")).readlines()]

# 模型名称
pretrained_weights = 'bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(pretrained_weights)
config = AutoConfig.from_pretrained(pretrained_weights, num_labels=len(class_list))
# 单独指定config，在config中指定分类个数
# 因为是分类任务，用 AutoModelForSequenceClassification
nlp_classif = AutoModelForSequenceClassification.from_pretrained(pretrained_weights,
                                                                 config=config)
# 指定机器
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 可能 gpu 显存不够
device = torch.device("cpu")
nlp_classif = nlp_classif.to(device)


time_start = time.time() #开始时间
epochs = 2
gradient_accumulation_steps = 1
max_grad_norm =0.1  #梯度剪辑的阀值

require_improvement = 1000                 # 若超过1000batch效果还没提升，则提前结束训练
savedir = './myfinetun-bert_chinese/'
os.makedirs(savedir, exist_ok=True)


def get_time_dif(start_time):
    """获取已使用时间"""
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))


def train(model, traindataloder, testdataloder):
    """
    开始训练
    :param model:
    :param traindataloder:
    :param testdataloder:
    :return:
    """
    start_time = time.time()
    # 在训练模式下，模型会启用如dropout和batch normalization这样的正则化技术。
    model.train()
    # 获取模型中所有可训练参数及其名称。这样可以方便地对不同类型的参数应用不同的优化策略
    param_optimizer = list(model.named_parameters())
    # 不需要权重衰减（weight decay/L2正则化）的参数名称部分。通常包括偏置项（bias）和LayerNorm层中的偏差与权重参数
    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
    # 创建了两个参数组，分别对需要和不需要权重衰减的参数应用不同的权重衰减率。第一组设置了0.01的权重衰减，第二组不进行权重衰减
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]
    # ：使用AdamW优化器初始化模型参数。AdamW是对Adam优化器的一个改进版本，它确保了权重衰减在梯度更新之前被正确应用。
    # 这里的lr表示学习率，设置为5e-5；eps是Adam算法中的一个稳定系数，设置为1e-8
    optimizer = AdamW(optimizer_grouped_parameters, lr=5e-5, eps=1e-8)

    # 创建一个线性学习率调度器，并带有预热阶段（warmup）。这里没有设置预热步数（num_warmup_steps），
    # 意味着没有预热阶段；num_training_steps 设置为整个训练过程中迭代的总步数，
    # 即训练数据加载器循环次数乘以轮数（epochs）。随着训练的进行，学习率将按照预先设定的方式逐渐降低，从而有助于模型收敛并防止过拟合。
    scheduler = get_linear_schedule_with_warmup(optimizer,
                                                num_warmup_steps=0, num_training_steps=len(traindataloder) * epochs)


    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升

    for epoch in range(epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, epochs))
        # sku_name代表文本序列，labels代表对应的类别标签。
        for i, (sku_name, labels) in enumerate(traindataloder):
            model.train()
            # 使用BERT分词器对文本序列进行编码，并根据需要补全至最大长度，同时将结果转换为PyTorch张量格式
            ids = tokenizer.batch_encode_plus(sku_name,
                                               #                max_length=model.config.max_position_embeddings,  #模型的配置文件中就是512，当有超过这个长度的会报错
                                               pad_to_max_length=True, return_tensors='pt')#没有return_tensors会返回list！！！！
            # 清零优化器中的梯度累计信息，准备进行新的反向传播过程
            optimizer.zero_grad()
            # 将标签数据从CPU转移到指定设备（如GPU）上，并去除可能存在的额外维度
            labels = labels.squeeze().to(device)
            # 将编码后的输入ID、标签和注意力掩码传入模型进行前向传播计算，得到损失和其他输出
            outputs = model(ids["input_ids"].to(device), labels=labels,
                            attention_mask=ids["attention_mask"].to(device))
            # 从模型返回的结果中提取损失值和logits（未归一化的预测概率）
            loss, logits = outputs[:2]
            # 如果设置了梯度累积步骤大于1，则需要将损失除以这个值，这样在多个小批次上累积更新梯度
            if gradient_accumulation_steps > 1:
                loss = loss / gradient_accumulation_steps
            # 计算梯度并反向传播到模型参数
            loss.backward()
            # 每经过gradient_accumulation_steps次迭代后
            if (i + 1) % gradient_accumulation_steps == 0:
                # 对模型所有参数的梯度进行裁剪，防止梯度过大导致训练不稳定
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

            optimizer.step()  # 应用梯度更新模型参数。
            scheduler.step()  # 更新学习率调度器，根据当前训练步数调整学习率
            model.zero_grad()  # 再次清零梯度，为下一批次的训练做准备

            '''评估模型在训练集和验证集上的性能，并根据验证集上的表现做出相应的决策：'''
            if total_batch % 100 == 0:
                # 每多少轮输出在训练集和验证集上的效果
                truelabel = labels.data.cpu()  # 真实类别
                predic = torch.argmax(logits,axis=1).data.cpu()  # 预测类别
                #                predic = torch.max(outputs.data, 1)[1].cpu()
                train_acc = metrics.accuracy_score(truelabel, predic)  # 比较
                dev_acc, dev_loss = evaluate(model, testdataloder)  # 计算验证集上的准确率和损失值
                if dev_loss < dev_best_loss:
                    '''比较当前验证集损失与历史最优验证集损失（dev_best_loss），如果当前损失更低，
                    则更新最优损失并保存模型至预设路径（savedir），同时记录最后一次改善的批次索引（last_improve）'''
                    dev_best_loss = dev_loss
                    model.save_pretrained(savedir)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                # 输出当前迭代次数、训练损失、训练准确率、验证损失、验证准确率以及已用时间
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'
                print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))
                model.train()
            total_batch += 1  # 增加累计批次计数器
            '''如果自上次验证集损失下降以来已经过去了超过指定数量的批次（这里是1000批次），
            并且在这期间验证集损失未再降低，则自动停止训练，打印提示信息，并跳出循环。'''
            if total_batch - last_improve > require_improvement:
                # 验证集loss超过1000batch没下降，结束训练
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break


def evaluate(model, testdataloder):
    model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for sku_name, labels in testdataloder:
            ids = tokenizer.batch_encode_plus( sku_name,
                                               #                max_length=model.config.max_position_embeddings,  #模型的配置文件中就是512，当有超过这个长度的会报错
                                               pad_to_max_length=True,return_tensors='pt')#没有return_tensors会返回list！！！！

            labels = labels.squeeze().to(device)
            outputs = model(ids["input_ids"].to(device), labels=labels,
                            attention_mask =ids["attention_mask"].to(device) )

            loss, logits = outputs[:2]
            loss_total += loss
            labels = labels.data.cpu().numpy()
            predic = torch.argmax(logits, axis=1).data.cpu().numpy()
            labels_all = np.append(labels_all, labels)
            predict_all = np.append(predict_all, predic)
    acc = metrics.accuracy_score(labels_all, predict_all)
    return acc, loss_total / len(testdataloder)


train(nlp_classif, traindataloder, testdataloder)

代码输出结果会生成一个文件夹：myfinetun-bert_chinese 里面存放的是模型，最后会生成一个 best 模型，我这里没跑完哈，所以结果不全