Resume training from specific checkpoint with Pytorch?

mbaske · January 7, 2021, 3:47pm

Hi - With Tensorflow, I was just editing the checkpoints textfile if I wanted to resume training from a specific checkpoint. What do I do when using Pytorch (checkpoint.pt looks wild…)? Thanks!

celion_unity · January 8, 2021, 5:17pm

Looks like we always load from “checkpoint.pt”, so you should just be able to copy (or symlink) the checkpoint you want to that name.

github.com

Unity-Technologies/ml-agents/blob/2335fc7cb22f88f2e266cc128c7d66d8858dc73f/ml-agents/mlagents/trainers/model_saver/torch_model_saver.py#L82-L83


      
          model_path = os.path.join(load_path, "checkpoint.pt")
          saved_state_dict = torch.load(model_path)

We have MLA-1517 in our internal tracker to provide a better interface for this on the command line. It should be a pretty straightforward PR (I think) if anybody reading this wants to try to implement it.