Search code examples
huggingface-transformers

saving finetuned model locally

I'm trying to understand how to save a fine-tuned model locally, instead of pushing it to the hub.

I've done some tutorials and at the last step of fine-tuning a model is running trainer.train() . And then the instruction is usually: trainer.push_to_hub

But what if I don't want to push to the hub? I want to save the model locally, and then later be able to load it from my own computer into future task so I can do inference without re-tuning.

How can I do that?

eg: Initially load a model from hugging face:

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()

Somehow save the new trained model locally, so that next time I can pass

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

  • You can use the save_model method:

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

    Or alternatively, the save_pretrained method:

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

    Then, when reloading your model, specify the path you saved to:

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