How to save and load nn.Parameter() in pytorch, such that one can continue optimisation over it?

Question

I know how to store and load nn.Model, but can not find how to make a checkpoint for nn.Parameter. I tried this version, but the optimizer is not changing the nn.Parameter value after restoring.

from torch import nn as nn
import torch
from torch.optim import Adam

alpha = torch.ones(10)
lr = 0.001
alpha = nn.Parameter(alpha)
print(alpha)
alpha_optimizer = Adam([alpha], lr=lr)

for i in range(10):
   alpha_loss = - alpha.mean()
   alpha_optimizer.zero_grad()
   alpha_loss.backward()
   alpha_optimizer.step()
   print(alpha)
path = "./test.pt"
state = dict(alpha_optimizer=alpha_optimizer.state_dict(),
               alpha=alpha)
torch.save(state, path)
checkpoint = torch.load(path)
alpha = checkpoint["alpha"]
alpha_optimizer.load_state_dict(checkpoint["alpha_optimizer"])
for i in range(10):
   alpha_loss = - alpha.mean()
   alpha_optimizer.zero_grad()
   alpha_loss.backward()
   alpha_optimizer.step()
   print(alpha)

Answer 1

The problem is that the optimizer is still with a reference to the old alpha (check id(alpha) vs. id(alpha_optimizer.param_groups[0]["params"][0]) before the last for loop), while a new one is set when you load it from the checkpoint in alpha = checkpoint["alpha"].

You need to update the params of the optimizer before loading its state:

# ....
torch.save(state, path)
checkpoint = torch.load(path)

# here's where the reference of alpha changes, and the source of the problem
alpha = checkpoint["alpha"]

# reset optim
alpha_optimizer = Adam([alpha], lr=lr)
alpha_optimizer.load_state_dict(checkpoint["alpha_optimizer"])

for i in range(10):
   alpha_loss = - alpha.mean()
   alpha_optimizer.zero_grad()
   alpha_loss.backward()
   alpha_optimizer.step()
   print(alpha)