【NLP，Huggingface，Colab】使用 Trainer 训练模型，并保存模型参数-Toy模板网

这篇具有很好参考价值的文章主要介绍了【NLP，Huggingface，Colab】使用 Trainer 训练模型，并保存模型参数。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

前置知识

Colab 的使用
Huggingface 官网和一些基础API

上代码

首先，建议保存代码到 VSCode，这样双击关键类，F12可以进入查看具体接口参数及其含义。
然后，建议代码在 Colab 运行，第一个是有默认 GPU 资源，第二个是否则会产生各种 ConnectionError, OSError 等错误…
重点可以看注释。自行摸索了一些额外的参数，大多数人都没有讲训练中/后需要保存模型参数…

"""
首先运行如下代码安装库
然后直接运行改代码即可
!pip install datasets transformers
!pip install accelerate -U
"""

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    DataCollatorWithPadding,
    TrainingArguments,
    AutoModelForSequenceClassification,
    Trainer,
)

# 加载数据集，并加载对应模型的分词器
raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)


def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)

# 数据集分词并打包，传给data_collator
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# 设置训练参数，这里我选择训练1poch，每处理20%steps就保存，注意最后100%时不保存。
training_args = TrainingArguments(
    "test-trainer",
    num_train_epochs=1,
    save_strategy="steps",
    save_steps=0.2,
)

# 设置模型
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

# 设置训练器，提供各种必要参数。
trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

# 训练，结束后保存模型
trainer.train()

model.save_pretrained("./output_model")