Saving and loading SGD optimizer

1

I’m trying to save SGD Optimizer state using
torch.save(optimizer.state_dict(), 'optim.pth') and then loading it with optimizer.load_state_dict(torch.load('optim.pth'))
But after loading the state, optimizer.step() gives error for in sgd.py for following line

if momentum != 0: param_state = self.state[p]
This is probably because it doesn't save the dictionary used by [optimizer](https://github.com/pytorch/pytorch/blob/master/torch/optim/optimizer.py) So I tried using `torch.save(optimizer, 'optim.pth')` but then the network simply doesn't train at all. So what' the correct way of saving and loading SGD optimizer (while using momentum) ?
2

Just found out that this is because it doesn’t save the dictionary used by [optimizer] (https://github.com/pytorch/pytorch/blob/master/torch/optim/optimizer.py) as it is while using optimizer.state_dict(). Before saving the optimizer, it’s state variable is defaultdict(<type 'dict'>, {}) but when it is loaded back it changes to a simple dict {}

3

So for now I’m training the network on a random batch before saving the state so that it doesn’t have to store a blank dict.

4

Could you, please explain then how to save and reload an optimizer.

I understand that if one is saving the model, the optimizer must also be saved otherwise the learning_rate value inside the optimizer is lost.

Thanks.

5

A quick workaround that fixed it for me was doing optimizer.state = defaultdict(dict, optimizer.state) after doing optimizer.load_state_dict(torch.load('optim.pth')).

2 Likes