Help needed on learning hyperparameter with SVI

Hi Pyro Lovers

I am attaching an example code for you to have a look. I wish to use SVI to train a GP with a hyperparameter for the kernel parameter variance which is a distribution. The training with SVI failed but the first block in pytorch training with retain_graph=True works okay. Any clues? Thanks

Second, if further I want to let the hyperparameter to follow another distribution with its parameters to be learned, how may I revise the code?

Thank you very much in advance.

J.

The code ==============================

import torch
import pyro
import pyro.distributions as dist
import pyro.contrib.gp as gp
from pyro.nn import PyroSample, PyroParam
from torch.distributions import constraints

clear the param store in case we’re in a REPL

pyro.clear_param_store()

Define the data

x = torch.linspace(-3, 3, 50)
y = torch.sin(x) + torch.randn(x.shape) * 0.2

Define the kernel function

kernel = gp.kernels.RBF(input_dim=1, variance=torch.tensor(0.5),
lengthscale=torch.tensor(1.0) )

Define a hyperparameter

kernel.lengthscale_hyper = PyroParam(torch.tensor(1.),
constraint=constraints.positive)
kernel.variance = PyroSample(dist.LogNormal(0, 1))
kernel.lengthscale = PyroSample(dist.LogNormal(0, kernel.lengthscale_hyper))
kernel.autoguide(“variance”, dist.Normal)
kernel.autoguide(“lengthscale”, dist.Normal)

Define the noise model

noise_param = torch.tensor(0.5)

Define the GP model with a probabilistic noise variance

gp_model = gp.models.GPRegression(x, y, kernel, noise=noise_param) # GPR will make this a parameter
gp_model.noise = PyroSample(dist.LogNormal(0, 1))
gp_model.autoguide(“noise”, dist.Normal)

Define the optimizer

optimizer = torch.optim.Adam(gp_model.parameters(), lr=0.005)
loss_fn = pyro.infer.Trace_ELBO().differentiable_loss
losses =
num_steps = 1000

The following block works

#for step in range(num_steps):

optimizer.zero_grad()

loss = loss_fn(gp_model.model, gp_model.guide)

loss.backward(retain_graph=True)

optimizer.step()

if step % 100 == 0:

print(“step {}: loss = {:.4f}”.format(step, loss))

losses.append(loss.item())

Define the inference procedure

optimizer = pyro.optim.Adam({“lr”: 0.01})
svi = pyro.infer.SVI(gp_model.model, gp_model.guide, optimizer, loss=pyro.infer.Trace_ELBO())

Train the model

for step in range(num_steps):
loss = svi.step()
if step % 100 == 0:
print(“step {}: loss = {:.4f}”.format(step, loss))

I found out a solution by setting retain_graph = True to Trace_ELBO()

J.