# Variational GP Regression with PyroΒΆ

[1]:

import math
import torch
import pyro
import gpytorch
from matplotlib import pyplot as plt

%matplotlib inline

[2]:

train_x = torch.linspace(0., 1., 21)
train_y = torch.pow(train_x, 2).mul_(3.7)
train_y = train_y.div_(train_y.max())
train_y += torch.randn_like(train_y).mul_(0.02)

fig, ax = plt.subplots(1, 1, figsize=(3, 2))
ax.plot(train_x.numpy(), train_y.numpy(), 'bo')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend(['Training data'])

[2]:

<matplotlib.legend.Legend at 0x7f61e1f8bc18>

[3]:

class PVGPRegressionModel(gpytorch.models.PyroGP):
def __init__(self, train_x, train_y, likelihood, name_prefix="mixture_gp"):
# Define all the variational stuff
variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
num_inducing_points=train_y.numel(),
)
variational_strategy = gpytorch.variational.VariationalStrategy(
self, train_x, variational_distribution
)

# Standard initializtation
super(PVGPRegressionModel, self).__init__(variational_strategy, likelihood, num_data=train_y.numel())
self.likelihood = likelihood

# Mean, covar
self.mean_module = gpytorch.means.ConstantMean()
self.covar_module = gpytorch.kernels.ScaleKernel(
gpytorch.kernels.MaternKernel(nu=1.5)
)

def forward(self, x):
mean = self.mean_module(x)  # Returns an n_data vec
covar = self.covar_module(x)
return gpytorch.distributions.MultivariateNormal(mean, covar)

[4]:

model = PVGPRegressionModel(train_x, train_y, gpytorch.likelihoods.GaussianLikelihood())

[5]:

from pyro import optim

def train(num_iter=400):
elbo = pyro.infer.Trace_ELBO(num_particles=256, vectorize_particles=True)
svi = pyro.infer.SVI(model.model, model.guide, optimizer, elbo)
model.train()

for i in range(num_iter):
loss = svi.step(train_x, train_y)
if not (i + 1) % 50:
print('Iter %d/%d - Loss: %.3f   lengthscale: %.3f   noise: %.3f' % (
i + 1, num_iter, loss,
model.covar_module.base_kernel.lengthscale.item(),
model.likelihood.noise.item()
))

%time train()

Iter 50/400 - Loss: 24.104   lengthscale: 0.643   noise: 0.474
Iter 100/400 - Loss: 19.331   lengthscale: 0.669   noise: 0.310
Iter 150/400 - Loss: 14.304   lengthscale: 0.709   noise: 0.193
Iter 200/400 - Loss: 6.317   lengthscale: 0.769   noise: 0.111
Iter 250/400 - Loss: -0.906   lengthscale: 0.815   noise: 0.056
Iter 300/400 - Loss: -9.677   lengthscale: 0.826   noise: 0.026
Iter 350/400 - Loss: -15.517   lengthscale: 0.855   noise: 0.013
Iter 400/400 - Loss: -19.656   lengthscale: 0.986   noise: 0.007
CPU times: user 21.1 s, sys: 28.5 s, total: 49.6 s
Wall time: 7.11 s

[6]:

fig, ax = plt.subplots(1, 1, figsize=(4, 3))
train_data, = ax.plot(train_x.cpu().numpy(), train_y.cpu().numpy(), 'bo')

model.eval()
output = model.likelihood(model(train_x))

mean = output.mean
lower, upper = output.confidence_region()
line, = ax.plot(train_x.cpu().numpy(), mean.detach().cpu().numpy())
ax.fill_between(train_x.cpu().numpy(), lower.detach().cpu().numpy(),
upper.detach().cpu().numpy(), color=line.get_color(), alpha=0.5)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend([train_data, line], ['Train data', 'Prediction'])

[6]:

<matplotlib.legend.Legend at 0x7f61de116320>

[ ]: