[NLP, Huggingface, Colab] Use Trainer to train the model and save the model parameters

Pre-knowledge

• Use of Colab
• Huggingface official website and some basic APIs

upper code

• First of all, it is recommended to save the code toVSCode , so double-click the key class, F12you can enter to view the specific interface parameters and their meanings.
  Then, it is suggested that the code is Colabrunning , the first one is to have the default GPUresource, and the second one will generate various ConnectionError, OSErroretc. errors...
• The focus can be seen in the notes. I explored some additional parameters by myself. Most people didn't talk about the need to save model parameters during/after training...

"""
首先运行如下代码安装库
然后直接运行改代码即可
!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")

• The final folder is as follows, test-trainer saves the training breakpoint, and output_model saves the parameter model after training.