Checkpointing optimizers about labml HOT 5 CLOSED

Got it! Ok, it's all I could desire about checkpointing ahah.
I'll close the issue.

from labml.

I also opened a pull request to show what I mean. Obviously you designed this and you'll know way better than me what's going on under the hood, so I may not see every problem that comes with my approach. What do you think?

from labml.

Hey, thanks for raising this and the pull request. Your approach makes perfect sense.

The decision to save to numpy was done way back when lab was for Tensorflow, and I left it when we added PyTorch support and later dropped Tensorflow support (because none of the lab users seemed to be using Tensorflow).

I have personally being resetting optimizers (to save checkpoint sizes and because it didn't have a huge effect on the tasks I was working on) so I didn't come across this problem yet. Saving as state_dict directly is both simpler and probably future proof to any changes PyTorch might make too.

I'll merge you request, ASAP. I just want to add some backward compatibility to load back old checkpoints saved in numpy formats. I will close the issue once we merge it.

Thanks again

from labml.

Great! In the checkpoint saving it could be needed to load also the current epoch/iteration. I think it could be saved in the .json file, but I don't know which could be the best practise to load that if we want to resume an interrupted training. Do you have anything in mind?

from labml.

Right now it picks the global step from the checkpoint. The checkpoint folder name is the step at which it was saved

from labml.