Bootstrap

一起学Hugging Face Transformers(10)- 使用Transformers 库的 Trainer API 进行模型微调


前言

在自然语言处理(NLP)领域,预训练模型如 BERT、GPT 等已经展示了其强大的能力。然而,实际应用中,预训练模型往往需要进一步微调(Fine-tuning)以适应具体任务。Hugging Face Transformers 库提供了强大的 Trainer API,使得模型微调变得简单高效。本文将详细介绍如何使用 Trainer API 对模型进行微调。

二、 环境准备

首先,确保你已经安装了 Hugging Face Transformers 库和其他必要的依赖:

pip install transformers datasets

此外,我们还需要安装 PyTorch 或 TensorFlow。本教程使用 PyTorch:

pip install torch

三、 数据准备

Hugging Face 提供了丰富的数据集接口。我们将使用 datasets 库加载一个示例数据集。这里以 SQuAD 数据集为例:

from datasets import load_dataset

dataset = load_dataset("squad")
train_dataset = dataset["train"]
eval_dataset = dataset["validation"]

四、 模型选择与加载

接下来,我们选择并加载一个预训练模型和对应的分词器。这里我们选择 distilbert-base-uncased

from transformers import AutoTokenizer, AutoModelForQuestionAnswering

model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

五、 数据预处理

我们需要对数据进行预处理,将问题和上下文编码成模型可以接受的输入格式:

def preprocess_function(examples):
    questions = [q.strip() for q in examples["question"]]
    inputs = tokenizer(
        questions,
        examples["context"],
        max_length=384,
        truncation="only_second",
        stride=128,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    sample_mapping = inputs.pop("overflow_to_sample_mapping")
    offset_mapping = inputs.pop("offset_mapping")

    start_positions = []
    end_positions = []

    for i, offset in enumerate(offset_mapping):
        input_ids = inputs["input_ids"][i]
        cls_index = input_ids.index(tokenizer.cls_token_id)

        sequence_ids = inputs.sequence_ids(i)

        sample_index = sample_mapping[i]
        answers = examples["answers"][sample_index]
        if len(answers["answer_start"]) == 0:
            start_positions.append(cls_index)
            end_positions.append(cls_index)
        else:
            start_char = answers["answer_start"][0]
            end_char = start_char + len(answers["text"][0])

            token_start_index = 0
            while sequence_ids[token_start_index] != 1:
                token_start_index += 1

            token_end_index = len(input_ids) - 1
            while sequence_ids[token_end_index] != 1:
                token_end_index -= 1

            if not (offset[token_start_index][0] <= start_char and offset[token_end_index][1] >= end_char):
                start_positions.append(cls_index)
                end_positions.append(cls_index)
            else:
                while token_start_index < len(offset) and offset[token_start_index][0] <= start_char:
                    token_start_index += 1
                start_positions.append(token_start_index - 1)

                while offset[token_end_index][1] >= end_char:
                    token_end_index -= 1
                end_positions.append(token_end_index + 1)

    inputs["start_positions"] = start_positions
    inputs["end_positions"] = end_positions
    return inputs

train_dataset = train_dataset.map(preprocess_function, batched=True, remove_columns=dataset["train"].column_names)
eval_dataset = eval_dataset.map(preprocess_function, batched=True, remove_columns=dataset["validation"].column_names)

六、 微调模型

现在,我们可以使用 Trainer API 进行模型微调。我们需要定义训练参数,并创建 Trainer 实例:

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="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(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
)

开始训练:

trainer.train()

七、 模型评估

训练完成后,我们可以对模型进行评估:

trainer.evaluate()

八、 保存模型

最后,我们可以保存微调后的模型:

model.save_pretrained("./finetuned_model")
tokenizer.save_pretrained("./finetuned_model")

总结

通过 Hugging Face Transformers 库的 Trainer API,我们可以方便快捷地对预训练模型进行微调。本文以 SQuAD 数据集为例,展示了从数据准备、模型选择、数据预处理到模型微调和评估的全过程。这种方法不仅简化了微调过程,还提高了训练效率和效果,适用于各种自然语言处理任务。

;