HuggingFace——Trainer的简单使用

最新推荐文章于 2025-09-06 14:46:22 发布

原创

最新推荐文章于 2025-09-06 14:46:22 发布 · 8k 阅读

27 ·

CC 4.0 BY-SA版权

文章标签：

#python #HuggingFace #Trainer #Pytorch

博客介绍了使用🤗Transformers的Trainer API微调模型的方法，包括定义类、模型、传入参数、评估设置和开始训练等步骤。还阐述了原生Pytorch训练方法，涉及定义优化器、调度器、训练位置等，最后介绍了用🤗 Accelerate库加速训练循环。

使用 Trainer API 微调模型[中文Course|API ]

🤗Transformers提供了一个 Trainer 类来帮助在数据集上微调任何预训练模型。

在定义Trainer之前首先要定义一个`TrainingArguments`类。

它将包含 Trainer用于训练和评估的所有超参数。其中唯一必须提供的参数是保存训练模型的目录——output_dir（ The output directory where the model predictions and checkpoints will be written.）参数。对于其余的参数，使用默认值。

定义模型

以分类句子模型为例，第二步是定义我们的模型。正如在将使用 AutoModelForSequenceClassification 类，它有两个参数：

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

在实例化此预训练模型后会收到警告。这是因为 BERT 没有在句子对分类方面进行过预训练，所以预训练模型的原来的头部(分类器或者说线性层)已经被丢弃，而是添加了一个适合句子序列分类的新头部。警告表明一些权重没有使用（对应于丢弃的预训练头的那些），而其他一些权重被随机初始化（新头的那些）。

传入参数

确定了模型之后，就可以定义一个Trainer通过将之前构造的所有对象传递给它——model、training_args，训练和验证数据集，data_collator，和tokenizer ：

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

评估设置

为了查看模型在每个训练周期结束的好坏，下面是使用**compute_metrics()**函数定义一个新的 Trainer。现在建立compute_metric()函数来较为直观地评估模型的好坏，可以使用 🤗 Evaluate 库中的指标。

def compute_metrics(eval_preds):
    metric = evaluate.load("glue", "mrpc") # 加载与 MRPC 数据集关联的指标
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels) # 返回的对象有一个 compute()方法我们可以用来进行度量计算的方法：

开始训练

只需要调用Trainer的train() 方法：

training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"