如何保存变压器 gpt2 的检查点以继续训练？

Question

我正在重新训练 GPT2 语言模型，并关注此博客：

https://towardsdatascience.com/train-gpt-2-in-your-own-language-fc6ad4d60171

在这里，他们训练了GPT2 上的网络，我正在尝试重新创建一个相同的网络。但是，我的数据集太大（250Mb），所以我想继续间隔训练。换句话说，我想检查模型训练。我怎么能这样做呢？

Answer 1

training_args = TrainingArguments(
    output_dir=model_checkpoint,
    # other hyper-params
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_set,
    eval_dataset=dev_set,
    tokenizer=tokenizer
)

trainer.train()
# Save the model to model_dir
trainer.save_model()

def prepare_model(tokenizer, model_name_path):
    model = AutoModelForCausalLM.from_pretrained(model_name_path)
    model.resize_token_embeddings(len(tokenizer))
    return model

# Assume tokenizer is defined, You can simply pass the saved model directory path.
model = prepare_model(tokenizer, model_checkpoint)