Can I save my own pretrained models and fine tune them later?

Hello,
I have a question regarding the save/load procedures described here:
https://pytorch.org/tutorials/beginner/saving_loading_models.html

I want to know if i am able to:
-Train a model on some data (lots of general domain data like movie reviews, product reviews, etc, that are associated with positive/negative classes)
-Save it
-Load the pre-trained model
-Train again with less data ( but task-specific like hate speech where the classes might be offensive/not-offensive ) and fine tune it by using a lower learning rate and try other techniques.

So, after I load my pre-trained model, can I start training again and modify it? With the procedures described in the link i mentioned above.

Thank you very much!

Yes, that would be a classic fine-tuning task and is possible in PyTorch.
As described in the docs you’ve posted, you might also need to save and load the optimizer’s state_dict, if your optimizer has internal states (e.g. Adam uses running estimates).

The Finetuning tutorial explains how to load pre-trained torchvision models and fine-tune them. Instead of loading the torchvision model, you might just want to load your custom model and adapt the code appropriately.

1 Like