[TOC]
transformers+huggingface训练模型
导入必要的库:
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
import numpy as np
import evaluate
CopyInsert
- 导入
datasets
用于加载数据集。 - 导入
transformers
中的组件,以便使用预训练的 BERT 模型和 tokenizer。 - 导入
numpy
用于数值计算。 - 导入
evaluate
用于计算模型预测的指标(这里是准确率)。
数据集下载:
dataset = load_dataset("yelp_review_full")
CopyInsert
- 从 Hugging Face 的数据集中下载 Yelp 评论数据集,该数据集包含各种评论和意见。
数据预处理:
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True)
CopyInsert
- 使用预训练的 BERT model("bert-base-cased")初始化 tokenizer。
- 定义
tokenize_function
函数,将评论文本编码成模型可接受的格式,设置填充和截断。
应用数据预处理:
tokenized_datasets = dataset.map(tokenize_function, batched=True)
CopyInsert
- 对下载的数据集应用
tokenize_function
,批量处理文本数据。
数据抽样:
small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))
CopyInsert
- 从训练和测试集中各随机抽取 1000 条样本,以加快训练速度和验证模型性能。
模型加载与训练配置:
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)
CopyInsert
- 加载预训练的 BERT 模型,并指定输出标签数(5个分类)。
model_dir = "models/bert-base-cased-finetune-yelp"
training_args = TrainingArguments(
output_dir=model_dir,
per_device_train_batch_size=16,
num_train_epochs=5,
logging_steps=100
)
CopyInsert
- 定义模型保存路径和训练参数,如每个设备的训练批大小、训练轮数和日志记录的频率。
指标评估:
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels)
CopyInsert
- 加载准确率评估指标。
- 定义
compute_metrics
函数,通过计算预测标签和真实标签的比较来评估模型性能。
实例化 Trainer:
trainer = Trainer(
model=model,
args=training_args,
train_dataset=small_train_dataset,
eval_dataset=small_eval_dataset,
compute_metrics=compute_metrics
)
CopyInsert
- 创建训练器
Trainer
的实例,用于处理模型的训练过程和评估。
开始训练:
trainer.train()
CopyInsert
- 运行训练过程。
监控 GPU 使用:
# 使用命令行工具: watch -n 1 nvidia-smi
CopyInsert
- 提供了一个命令行工具提示,以监控 GPU 的使用情况。
保存模型和训练状态:
trainer.save_model(model_dir)
trainer.save_state()
CopyInsert
- 保存训练完成后的模型和状态,以便后续使用。