How to resume training ?

[cmdrootaccess]

Hello,

how to resume training on a saved .pth ?

[akhil451]

if(model_name=='u2net'):
    net = U2NET(3, 1)
elif(model_name=='u2netp'):
    net = U2NETP(3,1)
net.load_state_dict(torch.load(saved_model_dir))

if torch.cuda.is_available():
    net.cuda()

After instantiating a model, you can load the .pth file after that.

[chenyangh]

@akhil451 Thanks for the answer. In addition, to resume the training from where exactly it was, one usually needs to save and load the optimizer (especially for Adam). https://pytorch.org/tutorials/beginner/saving_loading_models.html

[cmdrootaccess]

@chenyangh @akhil451 Thanks

[muhammadabdullah34907]

@chenyangh how about fine tuning the u2net.pth model by loading optimizer ? How to do load the optimizer for pretrained model ?

[chenyangh]

@muhammadabdullah34907 I don't think we have saved optimizer during the training. If you simply load the released pre-trained model and keep training it with a newly initialized Adam optimizer, it may diverge in the beginning.
In the case you want to train the model from scratch yourself, you can add under this line something like this:' torch.save(optimizer.state_dict(),, "some_file_name"))`. Just make sure that the model checkpoint and the optimizer states are matching. You can look at this tutorial: https://pytorch.org/tutorials/beginner/saving_loading_models.html for more details.

[cmdrootaccess]

@chenyangh how about saving the optimizer in one .pth together with the model state like this ?

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)