How to use the stop/resume training method within a custom training loop

[btetaren]

Hello,

is there a way to stop/resume training, as shown in the tutorial (i.e., loading the pickled density estimator from disc and continuing training using the resume_training=True in .train() ) but instead within a custom training loop like the one shown below?

density_estimator = build_nsf(dummy_theta,dummy_x,embedding_net=embedding_net)
optw = AdamW(list(density_estimator.parameters()), lr=5e-4)

num_epochs = 300
for ep in range(num_epochs):
for _, (theta_batch, x_batch) in enumerate(train_loader):
optw.zero_grad()
losses = density_estimator.loss(theta_batch, condition=x_batch)
loss = torch.mean(losses)
loss.backward()
optw.step()
if ep % 10 == 0:
print("last loss", loss.item())

If I pickle the density estimator, then reload it later from disc, how do I tell it to continue training (from where it left off) within the custom training loop?

Any help is much appreciated,

thanks.

[michaeldeistler]

Hey! With the custom training loop, everything is pytorch, so you can simply pickle (and then load) also the optimizer, and do not re-initialize the network.

Michael

[btetaren]

thank you!