HuggingFace Transformers教程：使用Trainer API微调模型

HuggingFace Transformers教程：使用Trainer API微调模型

course The Hugging Face course on Transformers 项目地址: https://gitcode.com/gh_mirrors/cou/course

前言

在自然语言处理领域，预训练模型的微调(fine-tuning)已成为标准实践。HuggingFace Transformers库提供的Trainer API极大简化了这一过程，使开发者能够专注于模型和任务本身，而不必纠结于训练循环的实现细节。本文将深入解析如何使用Trainer API高效地微调预训练模型。

准备工作

在开始微调前，我们需要完成以下准备工作：

1. 数据集加载：使用load_dataset方法加载目标数据集
2. 分词器初始化：选择与预训练模型匹配的分词器
3. 数据预处理：定义分词函数并应用到整个数据集

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding

# 加载MRPC数据集（句子对分类任务）
raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# 定义分词函数
def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)

# 应用分词函数到整个数据集
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

训练配置

Trainer API的核心是TrainingArguments类，它封装了所有训练相关的超参数：

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="model_output",  # 模型保存路径
    evaluation_strategy="epoch",  # 每个epoch结束后评估
    learning_rate=2e-5,  # 学习率
    per_device_train_batch_size=16,  # 每个设备的训练batch大小
    per_device_eval_batch_size=16,  # 每个设备的评估batch大小
    num_train_epochs=3,  # 训练epoch数
    weight_decay=0.01,  # 权重衰减
    save_strategy="epoch",  # 保存策略
    load_best_model_at_end=True,  # 训练结束时加载最佳模型
)

模型初始化

对于序列分类任务，我们需要使用AutoModelForSequenceClassification类：

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint, 
    num_labels=2  # MRPC是二分类任务
)

注意：当使用预训练模型进行新任务时，模型头部会被替换，因此会收到关于权重初始化的警告，这是正常现象。

构建Trainer

准备好所有组件后，我们可以实例化Trainer：

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

评估指标

为了在训练过程中监控模型性能，我们需要定义评估指标计算函数：

import numpy as np
import evaluate

def compute_metrics(eval_pred):
    metric = evaluate.load("glue", "mrpc")
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

然后将该函数添加到Trainer中：

trainer.compute_metrics = compute_metrics

开始训练

一切就绪后，启动训练过程：

trainer.train()

高级训练技巧

混合精度训练

通过启用FP16混合精度训练，可以显著减少显存占用并加速训练：

training_args.fp16 = True

梯度累积

当GPU显存有限时，可以使用梯度累积模拟更大的batch size：

training_args.gradient_accumulation_steps = 4  # 实际batch size = batch_size * 4

学习率调度

Trainer支持多种学习率调度策略：

training_args.lr_scheduler_type = "cosine"  # 余弦退火

常见问题解答

Q: 为什么我的评估指标没有显示？ A: 请确保设置了evaluation_strategy并提供了eval_dataset和compute_metrics函数。

Q: 如何恢复中断的训练？ A: 可以通过指定resume_from_checkpoint=True来从检查点恢复训练。

Q: 训练速度很慢怎么办？ A: 可以尝试启用混合精度训练(fp16=True)、增大batch size或使用梯度累积。

结语

通过本文，我们全面了解了如何使用HuggingFace Transformers的Trainer API进行模型微调。Trainer API封装了现代深度学习的最佳实践，使开发者能够专注于模型和任务本身，而无需从头实现训练循环。掌握这些技巧后，你可以高效地将预训练模型适配到各种NLP任务中。

course The Hugging Face course on Transformers 项目地址: https://gitcode.com/gh_mirrors/cou/course