Recommended approach to save and load optimiser state

[georgeherbert]

Hi, I was just wondering what the recommended approach is to save and load optimiser state?

I see for models (i.e. those that inherit from nn.Module) there are save_weights and load_weights methods. However, I don't see any analogous methods for the optimisers.

e.g. if I initialise optimiser = optim.AdamW(learning_rate=LEARNING_RATE) and train for n iterations, what's the best way to store the moment vectors so I can resume training from this point?

I started hacking around and implemented something of the form:

def load_optimiser(optimiser: optim.Optimizer, path: Path) -> None:
    params = list(mx.load(str(path)).items())
    optimiser.state = tree_unflatten(params)

However, I figured I'd ask if there's an existing approach built into the library that I just can't find before I go any further down this route. Many thanks.

[awni]

There is in fact no built-in saving/loading for optimizers, but I think what you implemented is a perfectly good way to do it.