Hugging Face
的 Trainer
类提供了一个高度集成的训练、评估和预测接口,能大幅简化我们预训练和微调的工作量。比如我最近最常会遇到的用监督学习微调语言模型任务为例,只需要定义一个Trainer,然后
传入我想要训练的模型和数据集,就可以直接运行微调任务。
以下是 Trainer
的一些常见用法和参数配置:
首先,Trainer
可以通过提供模型、数据集、训练参数等来进行快速的模型训练:
from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset# 加载预训练模型和分词器
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")# 加载数据集
dataset = load_dataset("glue", "sst2")
encoded_dataset = dataset.map(lambda x: tokenizer(x['sentence'], padding="max_length", truncation=True), batched=True)# 设置训练参数
training_args = TrainingArguments(output_dir="./results", # 模型保存路径evaluation_strategy="epoch", # 每个epoch进行一次评估learning_rate=2e-5, # 学习率per_device_train_batch_size=16, # 每个设备的训练批次大小per_device_eval_batch_size=16, # 每个设备的评估批次大小num_train_epochs=3, # 训练轮数weight_decay=0.01, # 权重衰减
)# 创建Trainer
trainer = Trainer(model=model, # 预训练模型args=training_args, # 训练参数train_dataset=encoded_dataset['train'], # 训练数据集eval_dataset=encoded_dataset['validation'], # 评估数据集
)# 开始训练
trainer.train()