NLP:微调BERT进行文本分类

本篇博客的重点在于BERT的使用。

transformers包版本:4.44.2

1. 微调BERT进行文本分类

  这里我们使用stanford大学的SST2数据集来演示BERT模型的微调过程。SST-2数据集(Stanford Sentiment Treebank 2)是一个用于情感分类的经典数据集,常用于自然语言处理(NLP)领域的情感分析任务。

  • 第1步: 下载数据。其代码如下:
python 复制代码
import pandas as pd
from transformers import BertTokenizer
from datasets import DatasetDict, Dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments

splits = {'train': 'data/train-00000-of-00001.parquet', 
          'validation': 'data/validation-00000-of-00001.parquet', 
          'test': 'data/test-00000-of-00001.parquet'}
train = pd.read_parquet("hf://datasets/stanfordnlp/sst2/" + splits["train"])
validation = pd.read_parquet("hf://datasets/stanfordnlp/sst2/" + splits["validation"])
test = pd.read_parquet("hf://datasets/stanfordnlp/sst2/" + splits["test"])
dataset = DatasetDict({'train': Dataset.from_pandas(train), 
                       'validation': Dataset.from_pandas(validation), 
                       'test': Dataset.from_pandas(test)})

要注意一下,这里并没有使用datasets包从hugging face上直接下载数据集的方式来获取数据,这是因为使用load_datesets方法获取数据时仍然会提示:NotImplementedError: Loading a dataset cached in a LocalFileSystem is not supported

  • 第2步: 构造训练集、验证集和测试集

SST2数据集中训练集(train)共67349条,验证集(validation)共872条,而测试集(test)共1821条。因为训练集数目较大微调会比较耗时,所以从这三个数据集分别抽取出了1000条、200条、200条进行后续的任务。具体代码如下:

python 复制代码
dataset['train'] = dataset['train'].shuffle(seed=42).select(range(1000))
dataset['validation'] = dataset['validation'].shuffle(seed=42).select(range(200))
dataset['test'] = dataset['test'].shuffle(seed=42).select(range(200))
print(dataset)

其输出结果如下:

bash 复制代码
Dataset({
    features: ['idx', 'sentence', 'label'],
    num_rows: 1000
})
Dataset({
    features: ['idx', 'sentence', 'label'],
    num_rows: 200
})
Dataset({
    features: ['idx', 'sentence', 'label'],
    num_rows: 200
})
  • 第3步:从bert中提取嵌入

训练集、验证集及测试集生成后,接着需要将这些语料全都转化成embedding向量。具体代码如下:

python 复制代码
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
def tokenize_function(examples):
    return tokenizer(examples['sentence'], padding='max_length', truncation=True)
dataset =dataset.map(tokenize_function, batched=True)
dataset=dataset.remove_columns(['sentence',"idx"])
dataset=dataset.rename_column("label","labels")
dataset.set_format("torch")
train_dataset=dataset['train']
eval_dataset=dataset['validation']
test_dataset=dataset['test']
  • 第4步:模型训练。 具体代码如下:
python 复制代码
model=BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
training_args = TrainingArguments(
    output_dir='results',         
    per_device_train_batch_size=8,  
    per_device_eval_batch_size=8,      
    num_train_epochs=1,
)
trainer = Trainer(
    model=model,                         
    args=training_args,                 
    train_dataset=train_dataset,        
    eval_dataset=eval_dataset,
)
trainer.train()
trainer.evaluate()
trainer.save_model("results")

关于上述代码,有以下几点需要说明:

  • 训练模型的选择: tranformers库中有多个分类模型,其中BertForSequenceClassification类适用于序列分类任务,比如情感分析和文本分类;而BertForTokenClassification类适用于token级的分类任务,比如命名实体识别。
  • TrainingArguments方法中的主要参数及其作用如下表所示:
参数名 作用
output_dir 指定模型和训练日志保存的记录;
num_train_epochs 设置训练的周期数(即遍历整个训练数据集的次数,指的是整个训练集将被遍历多少次以进行训练);
per_device_train_batch_size 设置每个设备(如GPU)上的训练批次大小,训练批次是指在一次训练迭代中,模型同时处理的数据样本数量;
per_device_eval_batch_size 设置每个设备上的评估批次大小;
logging_dir 指定训练日志的保存目录;
evaluation_strategy 设置评估策略。可以是 'no'(不评估)、'steps'(每隔一定步数评估)或 'epoch'(每个周期评估);
save_total_limit 设置保存模型检查点的总数限制,超过限制的检查点会被删除;
fp16 启用半精度浮点数(FP16)训练,以减少显存使用并加速训练(需要支持 FP16 的硬件);

参考资料

相关推荐
Solange__6 小时前
航天科普怎么用AI生成3D场景?从科学约束、尺度分层到网页展示的完整方案
图像处理·人工智能·游戏·3d·ai·自然语言处理·关卡设计
kyrie_sakura8 小时前
大模型学习18 大模型学习 - LangChain2_Model I/O
学习·自然语言处理
zhangfeng11338 小时前
大语言模型 Greedy Decoding,贪心解码 / 贪心搜索
人工智能·语言模型·自然语言处理
博谷10 小时前
什么是GEO监测?核心指标、工作原理与AI搜索诊断监测全解析
人工智能·机器学习·自然语言处理
ai大模型-探索者12 小时前
意图识别三选一(BERT微调、LLM+Prompt、Jev)
人工智能·prompt·bert·jev
a puppy slide 葱12 小时前
arXiv _ 2025 _ Gaussian Masked Autoencoders
论文阅读·人工智能·python·分类
FelixZhang02815 小时前
不懂金融或不会编程,如何开始量化?从策略搭建到回测验证
人工智能·python·算法·数学建模·均值算法·自然语言处理·金融
Solange__15 小时前
AI模型编辑器、Blender和3D Agent怎么选?按5类修改任务比较局部编辑、拓扑与自动化
人工智能·3d·自然语言处理·编辑器·音视频·blender·游戏策划
AI人工智能+2 天前
智能文档抽取系统通过“OCR+大模型”双引擎,实现从图像到结构化数据的精准转化,突破传统OCR仅识别文字的局限
自然语言处理·ocr·大语言模型·文档抽取
小豆庄2 天前
视频脱敏:云端识别与本地推理的边界实测
数据库·yolo·自然语言处理·视觉检测·视频·设计规范·纹理贴图