使用 Bert 做文本分类,利用 Trainer 框架实现 二分类,事半功倍
简介
使用 AutoModelForSequenceClassification
导入Bert 模型。
很多教程都会自定义 损失函数,然后手动实现参数更新。
但本文不想手动微调,故使用 transformers 的 Trainer
自动微调。
人生苦短,我用框架,不仅可保证微调出的模型的效果,而且还省时间。
导包
import evaluate
import numpy as np
from datasets import load_dataset
from transformers import (AutoTokenizer,AutoModelForSequenceClassification,
)import torch
from torch import nnimport os
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'# AG_News 英文分类数据集
# ds = load_dataset("fancyzhx/ag_news")## 中文分类数据集
ds = load_dataset("lansinuote/ChnSentiCorp")
数据集的详情如下:
DatasetDict({train: Dataset({features: ['text', 'label'],num_rows: 9600})validation: Dataset({features: ['text', 'label'],num_rows: 1200})test: Dataset({features: ['text', 'label'],num_rows: 1200})
})
ds["train"][0]
{'text': '选择珠江花园的原因就是方便,有电动扶梯直接到达海边,周围餐馆、食廊、商场、超市、摊位一应俱全。酒店装修一般,但还算整洁。 泳池在大堂的屋顶,因此很小,不过女儿倒是喜欢。 包的早餐是西式的,还算丰富。 服务吗,一般','label': 1}
加载 Bert 模型
model_name = "bert-base-chinese"tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True,
)bert = AutoModelForSequenceClassification.from_pretrained(model_name,trust_remote_code=True,num_labels=2,
)
如果你无法联网的话,使用本地huggingface模型:
bert = AutoModelForSequenceClassification.from_pretrained(model_name,trust_remote_code=True,revision="c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f",local_files_only=True,num_labels=2,
)
查看 bert 分类模型的网络结构:
bert
如上图所示,Bert 的分类模型:在原生的 Bert 模型后,加了一个Linear。
下述是数据集转换函数:
def tokenize_func(item):global tokenizertokenized_inputs = tokenizer(item["text"],max_length=512,truncation=True,)return tokenized_inputs
tokenized_datasets = ds.map(tokenize_func,batched=True,
)
tokenized_datasets
的详情如下所示:
DatasetDict({train: Dataset({features: ['text', 'label', 'input_ids', 'token_type_ids', 'attention_mask'],num_rows: 9600})validation: Dataset({features: ['text', 'label', 'input_ids', 'token_type_ids', 'attention_mask'],num_rows: 1200})test: Dataset({features: ['text', 'label', 'input_ids', 'token_type_ids', 'attention_mask'],num_rows: 1200})
})
Train
from transformers import TrainingArgumentsargs = TrainingArguments("ChnSentiCorp_text_cls",eval_steps=8,evaluation_strategy="steps",save_strategy="epoch",save_total_limit=3,learning_rate=2e-5,num_train_epochs=3,weight_decay=0.01,per_device_train_batch_size=32,per_device_eval_batch_size=16,logging_steps=8,save_safetensors=True,overwrite_output_dir=True,# load_best_model_at_end=True,
)
TrainingArguments 的参数解释点击查看下述文章:
LLM大模型之Trainer以及训练参数
from transformers import DataCollatorWithPaddingdata_collator = DataCollatorWithPadding(tokenizer=tokenizer)
from transformers import Trainertrainer = Trainer(model=bert,args=args,train_dataset=tokenized_datasets["train"],eval_dataset=tokenized_datasets["validation"],data_collator=data_collator,# compute_metrics=compute_metrics,tokenizer=tokenizer,
)
trainer.train()
训练过程,在终端可以看见,训练和验证的损失值变化。
如果安装了 wandb
,并且在系统环境变量中,进行了设置。
训练过程和评估过程的记录会自动上传到wandb中。
wandb
若你想使用 wandb,自行进行安装;个人强烈推荐,一劳永逸,这样就无需自己绘图展示模型的训练过程了。
在模型训练的过程,进入 wandb https://wandb.ai/home 看看模型的现在的训练的过程。
上图是在 wandb 网站看到的图,横轴是 epoch ,纵轴是 loss。
蓝色折线是在验证集上的损失,橙色折线是在训练集上的损失。
可以很直观的看到,在训练集上的loss 小于 在验证集上的 loss。
predict
训练完成的模型,使用 predict
方法,在测试集上预测。
predictions = trainer.predict(tokenized_datasets["test"])
preds = np.argmax(predictions.predictions, axis=-1)
preds
输出结果:
array([1, 0, 0, ..., 1, 1, 0])
预测结果评估
def eval_data(data):predictions = trainer.predict(data)preds = np.argmax(predictions.predictions, axis=-1)metric = evaluate.load("glue", "mrpc")return metric.compute(predictions=preds, references=predictions.label_ids)
eval_data(tokenized_datasets["test"])
输出结果:
{'accuracy': 0.9475, 'f1': 0.9478908188585607}
总结
总体上看,本文做了一下数据集的处理,大模型的微调过程、模型权重报错、日志记录,这些过程全部由 transformers 的 Trainer 自动进行。
用好 框架, 事半功倍。当然前提是已经掌握了基础的手动参数微调。
参考资料
- huggingface 使用 Trainer API 微调模型