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

背景

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

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

from transformers import AutoModelForSequenceClassification

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

第二步:设置训练超参

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

from transformers import AutoTokenizer

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

第四步:加载数据集

from datasets import load_dataset

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

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

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

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

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

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

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

第八步:初始化Trainer

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

第九步:开始训练

trainer.train()

总结:

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

你可能感兴趣的:(人工智能,pytorch,人工智能,自然语言处理,transformer,Train)