使用 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 模型，我这里没跑完哈，所以结果不全
![[答疑]张学友和Neal Ford的区别](https://img-blog.csdnimg.cn/img_convert/645c8520ee274d1379a38b9dccbc82f8.png)
