在本地保存微调模型

Question

我正在尝试了解如何在本地保存经过微调的模型，而不是将其推送到中心。

我已经完成了一些教程，微调模型的最后一步是运行 trainer.train()。然后指令通常是：trainer.push_to_hub

但是如果我不想推送到中心怎么办？我想将模型保存在本地，然后能够从我自己的计算机将其加载到未来的任务中，这样我就可以在不重新调整的情况下进行推理。

我该怎么做？

例如：最初从拥抱面加载模型：

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=5)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

trainer.train()

不知何故将新训练好的模型保存到本地，方便下次通过

model = 'some local directory where model and configs (?) got saved'

Answer 1

您可以使用 save_model method:

trainer.save_model("path/to/model")

或者，save_pretrained method:

model.save_pretrained("path/to/model")

然后，在重新加载模型时，指定保存到的路径：

AutoModelForSequenceClassification.from_pretrained("path/to/model")

在本地保存微调模型

saving finetuned model locally

huggingface-transformers