Saving checkpoints during training and loading about performer-pytorch HOT 3 CLOSED

@ylhsieh

Try this:

Save:

torch.save(model.state_dict(), '/content/model.pth')

checkpoint = {'state_dict': model.state_dict(),'optimizer' :optim.state_dict()}
torch.save(checkpoint, '/content/model_sd_opt.pth')

Load:

torch.load('/content/model_sd_opt.pth')

Although loading maybe tricky due to many reasons. So try to read the docs/try a few things.

Hope this is helpful.

from performer-pytorch.

Hi thanks so much for your help! I end up using deepspeed and the following commands

model_engine.save('checkpoint_dir', f'{epoch}_{step}')
model_engine.load('checkpoint_dir')

from performer-pytorch.

@ylhsieh You are welcome. I am glad you found a working solution. Yes, deepspeed is a good way to go too, so whatever works :)

from performer-pytorch.