Applying Bayesian inference on multi Tasks exact GP

[lovecodng]

Hello,
I am trying to apply fully Bayesian inference on multitasks exact GP model. I have developed the following code. In the provided code I assumed that all hyperparameters should be learned, and I am doing this by integrating Pyro with gpytorch and run MCMC using nuts. I am new to gpytorch and Pyro so I am wondering If I am applying the concept correctly. I really appreciate any advice or hint on how to make the code run correctly without error.
Thank you in advance.

import math
import torch
import gpytorch
import pyro
from pyro.infer.mcmc import NUTS, MCMC, HMC
from matplotlib import pyplot as plt

%matplotlib inline
%load_ext autoreload
%autoreload 2

import pyro.distributions as dist
from gpytorch.priors import NormalPrior

from pyro.nn.module import PyroModule, pyro_method
from gpytorch.mlls import ExactMarginalLogLikelihood

# Generate data for training
train_x = torch.linspace(0, 1, 20)
train_y1 = train_x * 2 + 5
train_y2 = train_x * 4 + 10
train_y = torch.stack([train_y1, train_y2], -1)
noise_factor = 0.1
noise = noise_factor * torch.randn_like(train_y)
train_y = train_y + noise

class LinearMean(torch.nn.Module):
    def __init__(self, num_tasks):
        super().__init__()
        self.weights = torch.nn.Parameter(torch.abs(torch.randn(num_tasks)))
        self.bias = torch.nn.Parameter(torch.abs(torch.randn(1, num_tasks)))
        self.register_buffer("initialized", torch.tensor(0, dtype=torch.uint8))
        self.weights_prior = NormalPrior(0, 1)
        self.bias_prior = NormalPrior(0, 1)

    def _set_weights(self, value):
        self.weights.data = value

    def _set_bias(self, value):
        self.bias.data = value

    def forward(self, x):
        return x.matmul(self.weights.unsqueeze(-1).T) + self.bias

    def forward_with_nograd(self, x):
        with torch.no_grad():
            res = x.matmul(self.weights.unsqueeze(-1).T) + self.bias
        return res

    def initialize(self, train_x):
        if not self.initialized:
            self.weights_prior.loc = self.weights.mean().detach()
            self.weights_prior.scale = torch.tensor([1e-2])
            self.bias_prior.loc = self.bias.mean().detach()
            self.bias_prior.scale = torch.tensor([1e-2])
            self.initialized.fill_(1)


class MultitaskGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.train_inputs = train_x
        self.train_targets = train_y
        self.likelihood = likelihood
        num_tasks = train_y.shape[-1]
        self.mean_module = LinearMean(num_tasks)
        
        self.covar_module = gpytorch.kernels.MultitaskKernel(
            gpytorch.kernels.RBFKernel(lengthscale_prior=NormalPrior(0, 1)),
            num_tasks=num_tasks,rank=1,
        )
        self.covar_module.task_covar_module.outputscale = torch.nn.Parameter(
            torch.ones(num_tasks),
            requires_grad=True,
        )
        if not hasattr(self.covar_module.task_covar_module, "lengthscale"):
            self.covar_module.task_covar_module.lengthscale = torch.nn.Parameter(
                torch.ones(num_tasks),
                requires_grad=True,
            )

        self._set_outputscale = lambda value: self.covar_module.task_covar_module._set_outputscale(value)
        self._set_lengthscale = lambda value: self.covar_module.task_covar_module._set_lengthscale(value)
        self._set_bias = lambda value: self.mean_module._set_bias(value)
        self._set_weights = lambda value: self.mean_module._set_weights(value)

        self.register_buffer("train_x_buf", self.train_inputs)
        self.register_buffer("train_y_buf", self.train_targets)
        self.register_buffer("train_y_noise", torch.ones_like(self.train_targets))

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)

    def get_weights(self):
        return self.mean_module.weights.squeeze(-1).detach().cpu().numpy()

    def get_bias(self):
        return self.mean_module.bias.item()

    def get_outputscale(self):
        return self.covar_module.task_covar_module.outputscale.item()

    def get_lengthscale(self):
        return self.covar_module.task_covar_module.base_kernel.lengthscale.item()

    def initialize(self):
        self.mean_module.initialize(self.train_x_buf)
        self.covar_module.initialize()
        self.covar_module.task_covar_module.initialize()

        # set outputscale and lengthscale for each task
        num_tasks = self.train_targets.shape[-1]
        self.covar_module.task_covar_module.outputscale_constraint = gpytorch.constraints.Positive()
        self.covar_module.task_covar_module.lengthscale_constraint = gpytorch.constraints.Positive()
        # self.covar_module.task_covar_module.lengthscale = torch.nn.Parameter(torch.ones(num_tasks))
        # self.covar_module.task_covar_module.initialize(lengthscale=torch.ones(num_tasks))
        # self.covar_module.task_covar_module.initialize(lengthscale_prior=NormalPrior(0, 1))
        
        # set training data and likelihood noise
        # self.set_train_data(self.train_x_buf, self.train_y_buf
        print(self.train_y_noise.mean(dim=0).shape)
        self.likelihood.initialize(raw_noise=torch.zeros(num_tasks), noise=self.train_y_noise.mean(dim=0))

    def set_train_data(self, train_x, train_y):
        # Add an extra dimension to train_x
        train_x = train_x.unsqueeze(-1)
        self.train_inputs = train_x
        self.train_targets = train_y
        self.train_x_buf = train_x
        self.train_y_buf = train_y
        self.covar_module.set_train_data(train_x, train_y)
        
    def state_dict(self, *args, **kwargs):
        state_dict = super().state_dict(*args, **kwargs)
        state_dict["train_x_buf"] = self.train_x_buf
        state_dict["train_y_buf"] = self.train_y_buf
        state_dict["mean_module"] = self.mean_module.state_dict(*args, **kwargs)
        state_dict["covar_module"] = self.covar_module.state_dict(*args, **kwargs)
        state_dict["likelihood"] = self.likelihood.state_dict(*args, **kwargs)
        return state_dict

    def load_state_dict(self, state_dict, *args, **kwargs):
        self.train_x_buf = state_dict.pop("train_x_buf")
        self.train_y_buf = state_dict.pop("train_y_buf")
        self.mean_module.load_state_dict(state_dict.pop("mean_module"), *args, **kwargs)
        self.covar_module.load_state_dict(state_dict.pop("covar_module"), *args, **kwargs)
        self.likelihood.load_state_dict(state_dict.pop("likelihood"), *args, **kwargs)
        super().load_state_dict(state_dict, *args, **kwargs)
        
    @pyro_method
    def model(self):
        self.initialize()
        prior_dict = {
            'covar_module.task_covar_module.outputscale': NormalPrior(0, 1),
            'covar_module.task_covar_module.base_kernel.lengthscale': NormalPrior(0, 1),
            'mean_module.bias': NormalPrior(0, 1),
            'mean_module.weights': NormalPrior(0, 1)
        }
        pyro.module('gp', self)
        for name, prior in prior_dict.items():
            pyro.sample(name, prior.expand([1]))
        likelihood = pyro.module('likelihood', self.likelihood)
        output = self(self.train_x_buf)
        return likelihood(output, self.train_y_buf)

# Define likelihood and model
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(
    num_tasks=2,
    noise_prior=gpytorch.priors.NormalPrior(0, 1),
    noise_constraint=gpytorch.constraints.LessThan(1000000),
)

model = MultitaskGPModel(train_x, train_y, likelihood)

mll = ExactMarginalLogLikelihood(model.likelihood, model)

def model(train_x, train_y):
    likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=train_y.shape[1])
    model = MultitaskGPModel(train_x, train_y, likelihood)
    model.initialize()

    return model

def likelihood_fn(train_x, train_y, model):
    likelihood = model.likelihood
    with gpytorch.settings.fast_pred_var(), torch.no_grad():
        model.eval()
        observed_pred = likelihood(model(train_x))
        loss = -observed_pred.log_prob(train_y).sum()
    return loss

# Define the number of samples and burn-in
num_samples = 1000
warmup_steps = 1000

# Define the Pyro model
model = model(train_x, train_y)

# Define the likelihood function
likelihood_fn = lambda model: likelihood_fn(train_x, train_y, model)

# Define the NUTS sampler
nuts_kernel = NUTS(likelihood_fn)

# Define the MCMC sampler
mcmc = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps)

# Run the MCMC sampler
mcmc.run(model)

running the above code, gives the following error

RuntimeError                              Traceback (most recent call last)
File ~\anaconda3\envs\python3-env\lib\site-packages\gpytorch\module.py:110, in Module.initialize(self, **kwargs)
    109 try:
--> 110     self.__getattr__(name).data.copy_(val.expand_as(self.__getattr__(name)))
    111 except RuntimeError:

RuntimeError: The expanded size of the tensor (1) must match the existing size (2) at non-singleton dimension 0.  Target sizes: [1].  Tensor sizes: [2]

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
Cell In[93], line 21
     18 warmup_steps = 1000
     20 # Define the Pyro model
---> 21 model = model(train_x, train_y)
     23 # Define the likelihood function
     24 likelihood_fn = lambda model: likelihood_fn(train_x, train_y, model)

Cell In[93], line 4, in model(train_x, train_y)
      2 likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=train_y.shape[1])
      3 model = MultitaskGPModel(train_x, train_y, likelihood)
----> 4 model.initialize()
      6 return model

Cell In[92], line 124, in MultitaskGPModel.initialize(self)
    117 # self.covar_module.task_covar_module.lengthscale = torch.nn.Parameter(torch.ones(num_tasks))
    118 # self.covar_module.task_covar_module.initialize(lengthscale=torch.ones(num_tasks))
    119 # self.covar_module.task_covar_module.initialize(lengthscale_prior=NormalPrior(0, 1))
    120 
    121 # set training data and likelihood noise
    122 # self.set_train_data(self.train_x_buf, self.train_y_buf
    123 print(self.train_y_noise.mean(dim=0).shape)
--> 124 self.likelihood.initialize(raw_noise=torch.zeros(num_tasks), noise=self.train_y_noise.mean(dim=0))

File ~\anaconda3\envs\python3-env\lib\site-packages\gpytorch\module.py:115, in Module.initialize(self, **kwargs)
    113             self.__getattr__(name).data = val
    114         else:
--> 115             self.__getattr__(name).data.copy_(val.view_as(self.__getattr__(name)))
    117 elif isinstance(val, float):
    118     constraint = self.constraint_for_parameter_name(name)

RuntimeError: shape '[1]' is invalid for input of size 2

[Balandat]

This doesn't answer your specific question, but we have an implementation of a fully Bayesian MTGP in botorch: https://github.com/pytorch/botorch/blob/main/botorch/models/fully_bayesian_multitask.py

Note that that uses a pure pytorch model for fitting, from which we then load hyperparameter samples into a gpytorch model (this is for performance reasons since running the MCMC chain via pyro on the gpytorch model can be very slow).