入门微调预训练Transformer模型

大家好，HuggingFace 为众多开源的自然语言处理（NLP）模型提供了强大的支持平台，让这些模型能够通过训练和微调来更好地服务于各种特定的应用场景。在大型语言模型（LLM）迅猛发展的今天，HuggingFace 提供的核心工具，特别是 Trainer 类，极大地优化了 NLP 模型的训练过程，开发者得以更加高效地实现模型定制和优化。

HuggingFace 的 Trainer 类是为 Transformer 模型量身打造的，不仅优化了模型的交互体验，还与 Datasets 和 Evaluate 等库实现了紧密集成，支持更高级的分布式训练，并能无缝对接 Amazon SageMaker 等基础设施服务。通过这种方式，可以更加便捷地进行模型训练和部署。

本文将通过一个实例，展示如何利用 HuggingFace 的 Trainer 类在本地环境中对 BERT 模型进行微调，以处理文本分类任务。并且重点介绍如何使用 HuggingFace 模型中心的预训练模型，而不是深入机器学习的理论基础。

1.设置

示例将在 SageMaker Studio（https://aws.amazon.com/cn/sagemaker/studio/）环境下进行操作，利用 ml.g4dn.12xlarge 实例搭载的 conda_python3 内核来完成任务。需要提醒的是，可以选择使用更小型的实例，但这可能会影响训练速度，具体取决于可用的 CPU/工作进程的数量。

使用 HuggingFace 数据集库下载数据集。

import datasets
from datasets import load_dataset

这里指定了训练数据集和评估数据集，会在训练循环中进行使用。

train_dataset = load_dataset("imdb", split="train")
test_dataset = load_dataset("imdb", split="test")
test_subset = test_dataset.select(range(100)) # 取数据的一个子集进行评估

对于任何文本数据，必须指定一个分词器，将数据预处理成模型可以理解的格式。在这种情况下，这里指定了我们使用的 BERT 模型的 HuggingFace 模型中心 ID。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

# 分词文本数据
def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

然后使用内置的 map 函数处理我们的训练和评估数据集。

tokenized_train = train_dataset.map(tokenize_function, batched=True)
tokenized_test = test_subset.map(tokenize_function, batched=True)

预处理后的数据

2.微调 BERT

数据准备就绪后，利用先前选定的模型ID来加载BERT模型。需要注意的是，针对文本分类任务，还定义了标签的总数。在此案例中设定了两个标签，分别用0和1来表示，0代表负面，1代表正面。

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

接下来在训练循环中，需要定义一个TrainingArguments对象。在这个对象中，可以设置训练过程中的各种参数，比如训练周期的数量、分布式训练的策略等。

from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch", num_train_epochs=1)

对于评估，使用 Evaluate 库内置的评估函数。

import numpy as np
import evaluate
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)

然后将 TrainingArguments、分词数据集和评估指标函数传递给 Trainer 对象。可以使用 train 方法启动训练运行，这将需要大约 10-15 分钟的时间，具体取决于现有硬件。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test, # 使用测试作为评估
    compute_metrics=compute_metrics,
    tokenizer=tokenizer
)
trainer.train()

训练完成

对于推理，可以直接使用微调后的 trainer 对象，并在用于评估的分词测试数据集上进行预测：

trainer.predict(tokenized_test)

输出

在更为实际的应用场景中，可以使用 trainer 对象将模型工件保存到本地目录中。

trainer.save_model("./custom_model")

模型工件

然后可以加载这些模型工件，指定训练的模型类型，并在单个数据点上进行推理。

loaded_model = AutoModelForSequenceClassification.from_pretrained(pretrained_model_name_or_path="custom_model/")

# 样本推理
encoding = tokenizer("I am super delighted", return_tensors="pt")
res = loaded_model(**encoding)
predicted_label_classes = res.logits.argmax(-1)
predicted_label_classes