Saving and loading optimizer state

I am trying to load and fine tune Google’s ViT model using hugging face. I’m trying to use torch.optim.AdamW for optimisation.

How do I load and save state_dict of the optimizer that was defined in configure_optimizers callback?

Thank you.

Hi @thevishnupradeep, model state_dict is available with lit_model.state_dict() method. You can save it using and load it back with torch.load(...)

Here is an example -

class BoringModel(pl.LightningModule):
    def __init__(self):
        super().__init__() = 1e-3
        self.model = NeuralNet()
    def forward(self, x):
        return self.model(x)
    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters())

model = BoringModel()

# state available

# save the state, "state.pkl")

# load the state

Also, we have deprecated this forum so I would request you to use our Github Discussions for quicker response. :slight_smile: