在PyTorch里面利用transformers的Trainer微调预训练大模型

若石之上2023-09-21 10:31

背景

transformers提供了非常便捷的api来进行大模型的微调，下面就讲一讲利用Trainer来微调大模型的步骤

第一步：加载预训练的大模型

python 复制代码

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

第二步：设置训练超参

python 复制代码

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="path/to/save/folder/",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=2,
)

比如这个里面设置了epoch等于2

第三步：获取分词器tokenizer

python 复制代码

from transformers import AutoTokenizer

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

第四步：加载数据集

python 复制代码

from datasets import load_dataset

dataset = load_dataset("rotten_tomatoes")  # doctest: +IGNORE_RESULT

第五步：创建一个分词函数，指定数据集需要进行分词的字段：

python 复制代码

def tokenize_dataset(dataset):
    return tokenizer(dataset["text"])

第六步：调用map()来将该分词函数应用于整个数据集

python 复制代码

dataset = dataset.map(tokenize_dataset, batched=True)

第七步：使用DataCollatorWithPadding来批量填充数据，加速填充过程：

python 复制代码

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

第八步：初始化Trainer

python 复制代码

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
)  # doctest: +SKIP

第九步：开始训练

python 复制代码

trainer.train()

总结：

利用Trainer提供的api，只需要简简单单的九步，十几行代码就能进行大模型的微调，你要不要动手试一试？