[Documentation/Examples] qNEI with Deep Gaussian Process

[super145r]

Hello
I am trying to implement the qNEI with Deep gaussian process.
(https://docs.gpytorch.ai/en/v1.2.1/examples/05_Deep_Gaussian_Processes/Deep_Gaussian_Processes.html)
But this is quite difficult for me in this moment. So I am really appreciate if you show some code example for qNEI with DGP.
Thank you for considering my request.

[Balandat]

Hi @super145r, we have some prototypes for using deep GPs for BO, but iirc these aren't actually hooking things up into the Model API. Do you have any specific difficulties, or is it just that you don't know how to start?

[super145r]

Thanks a lot for your reply @Balandat
I don't know if my understanding is correct, but I implemented my test code. And it seems not to work properly.
The questions are as follows.

1. In this case, is it correct to use the Deterministic model?
2. Is qNEI (test_x.unsqueeze (-2)) correct? (at line 186)
3. If I run qNEI(train_x.unsqueeze (-2)), qNEI must be 0?

Here is my test code:

import torch
import tqdm
import gpytorch
from torch.nn import Linear
from gpytorch.means import ConstantMean, LinearMean
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.variational import VariationalStrategy, CholeskyVariationalDistribution
from gpytorch.distributions import MultivariateNormal
from gpytorch.models import ApproximateGP, GP
from gpytorch.mlls import VariationalELBO, AddedLossTerm
from gpytorch.likelihoods import GaussianLikelihood

from gpytorch.models.deep_gps import DeepGPLayer, DeepGP
from gpytorch.mlls import DeepApproximateMLL

from botorch.test_functions.synthetic import Branin

train_n=10
train_x=torch.rand((train_n,2))
testfunction=Branin()
bounds=testfunction._bounds
train_x[:,0]=(train_x[:,0])*(bounds[0][1]-bounds[0][0])+bounds[0][0]
train_x[:,1]=(train_x[:,1])*(bounds[1][1]-bounds[1][0])+bounds[1][0]
train_y=testfunction.evaluate_true(train_x)

test_x=torch.rand((50,2))
test_x[:,0]=(test_x[:,0])*(bounds[0][1]-bounds[0][0])+bounds[0][0]
test_x[:,1]=(test_x[:,1])*(bounds[1][1]-bounds[1][0])+bounds[1][0]
test_y=testfunction.evaluate_true(test_x)

from torch.utils.data import TensorDataset, DataLoader
train_dataset = TensorDataset(train_x, train_y)
train_loader = DataLoader(train_dataset, batch_size=1024, shuffle=True)

class ToyDeepGPHiddenLayer(DeepGPLayer):
    def __init__(self, input_dims, output_dims, num_inducing=128, mean_type='constant'):
        if output_dims is None:
            inducing_points = torch.randn(num_inducing, input_dims)
            batch_shape = torch.Size([])
        else:
            inducing_points = torch.randn(output_dims, num_inducing, input_dims)
            batch_shape = torch.Size([output_dims])

        variational_distribution = CholeskyVariationalDistribution(
            num_inducing_points=num_inducing,
            batch_shape=batch_shape
        )

        variational_strategy = VariationalStrategy(
            self,
            inducing_points,
            variational_distribution,
            learn_inducing_locations=True
        )

        super(ToyDeepGPHiddenLayer, self).__init__(variational_strategy, input_dims, output_dims)

        if mean_type == 'constant':
            self.mean_module = ConstantMean(batch_shape=batch_shape)
        else:
            self.mean_module = LinearMean(input_dims)
        self.covar_module = ScaleKernel(
            RBFKernel(batch_shape=batch_shape, ard_num_dims=input_dims),
            batch_shape=batch_shape, ard_num_dims=None
        )

        self.linear_layer = Linear(input_dims, 1)

    def forward(self, x):
        mean_x = self.mean_module(x) # self.linear_layer(x).squeeze(-1)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)

    def __call__(self, x, *other_inputs, **kwargs):
        """
        Overriding __call__ isn't strictly necessary, but it lets us add concatenation based skip connections
        easily. For example, hidden_layer2(hidden_layer1_outputs, inputs) will pass the concatenation of the first
        hidden layer's outputs and the input data to hidden_layer2.
        """
        if len(other_inputs):
            if isinstance(x, gpytorch.distributions.MultitaskMultivariateNormal):
                x = x.rsample()

            processed_inputs = [
                inp.unsqueeze(0).expand(self.num_samples, *inp.shape)
                for inp in other_inputs
            ]

            x = torch.cat([x] + processed_inputs, dim=-1)

        return super().__call__(x, are_samples=bool(len(other_inputs)))
    
num_output_dims = 10

from botorch.models.deterministic import GenericDeterministicModel

class qNEI_DeepGP(GenericDeterministicModel):
    def forward(self,x):
        with gpytorch.settings.num_likelihood_samples(1):
            prdictive_value=self._f.likelihood(self._f(x))
            
        return prdictive_value.mean.squeeze(0)

class DeepGP(DeepGP):
    def __init__(self, train_x_shape):
        hidden_layer = ToyDeepGPHiddenLayer(
            input_dims=train_x_shape[-1],
            output_dims=num_output_dims,
            mean_type='linear',
        )

        last_layer = ToyDeepGPHiddenLayer(
            input_dims=hidden_layer.output_dims,
            output_dims=None,
            mean_type='constant',
        )

        super().__init__()

        self.hidden_layer = hidden_layer
        self.last_layer = last_layer
        self.likelihood = GaussianLikelihood()

    def forward(self, inputs):
        hidden_rep1 = self.hidden_layer(inputs)
        output = self.last_layer(hidden_rep1)
        return output

    def predict(self, test_loader):
        with torch.no_grad():
            mus = []
            variances = []
            lls = []
            for x_batch, y_batch in test_loader:
                preds = model.likelihood(model(x_batch))
                mus.append(preds.mean)
                variances.append(preds.variance)
                lls.append(model.likelihood.log_marginal(y_batch, model(x_batch)))

        return torch.cat(mus, dim=-1), torch.cat(variances, dim=-1), torch.cat(lls, dim=-1)

model = DeepGP(train_x.shape)

num_epochs = 100
num_samples = 10


optimizer = torch.optim.Adam([
    {'params': model.parameters()},
], lr=0.01)
mll = DeepApproximateMLL(VariationalELBO(model.likelihood, model, train_x.shape[-2]))


for i in range(num_epochs):
    # Within each iteration, we will go over each minibatch of data
    minibatch_iter = train_loader
    for x_batch, y_batch in minibatch_iter:
        with gpytorch.settings.num_likelihood_samples(num_samples):
            optimizer.zero_grad()
            output = model(x_batch)
            loss = -mll(output, y_batch)
            loss.backward()
            optimizer.step()

                       
import gpytorch
import math


test_dataset = TensorDataset(test_x, test_y)
test_loader = DataLoader(test_dataset, batch_size=1024)

model.eval()
predictive_means, predictive_variances, test_lls = model.predict(test_loader)

rmse = torch.mean(torch.pow(predictive_means.mean(0) - test_y, 2)).sqrt()

from botorch.acquisition.monte_carlo import qExpectedImprovement, qNoisyExpectedImprovement
from botorch.sampling.samplers import SobolQMCNormalSampler

sampler = SobolQMCNormalSampler(256)
qneimodel=qNEI_DeepGP(model,num_outputs=1)
qNEI = qNoisyExpectedImprovement(qneimodel, train_x, sampler)
qneivalues = qNEI(test_x.unsqueeze(-2))


print(f"RMSE: {rmse.item()}, NLL: {-test_lls.mean().item()}")