本文主要是介绍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 timedeltadata_dir = 'THUCNews/data'
# 在代码开始部分添加全局变量和函数
global_batch_size = 4 # 初始化为固定的batch_size
max_batch_size = 32 # 根据实际情况设置最大允许的batch_sizedef 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_bsreturn global_batch_sizedef read_file(path):with open(path, 'r', encoding="UTF-8") as file:docus = file.readlines()newDocus = []for data in docus:newDocus.append(data)return newDocusclass Label_Dataset(dataset.Dataset): # 建立自定义数据集def __init__(self, data):self.data = datadef __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_timereturn 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-8optimizer = 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 # 记录进行到多少batchdev_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_lossmodel.save_pretrained(savedir)improve = '*'last_improve = total_batchelse: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 = Truebreakif flag:breakdef evaluate(model, testdataloder):model.eval()loss_total = 0predict_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 += losslabels = 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 模型，我这里没跑完哈，所以结果不全
这篇关于bert新闻标题分类的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
