[Bug] New release initialize MultitaskGaussianLikelihood

🐛 Bug

Hi and thanks for the great work.

Before the new release I used to initialize the value of a MultitaskGaussianLikelihood with

likelihood.noise_covar.noise = torch.tensor([0.04])
likelihood.noise = torch.tensor([1e-4])

but now this throws an error saying likelihood.noise_covar.noise does not exist and likelihood.noise is the wrong size for num_tasks > 1. Any idea how I am supposed to set the value of the MultitaskGaussianLikelihood now? See #1303 for the post that advised me to initialize this way in the first place.

Thanks!

To reproduce

import torch
import gpytorch
import logging
import math
import numpy as np
from matplotlib import pyplot as plt

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, kernel):
        train_x = torch.squeeze(train_x)
        train_y = torch.squeeze(train_y)
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = kernel

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

    def optimize(self, likelihood, train_x, train_y, training_iter=50,
                 optimizer='Adam'):
        train_x = torch.squeeze(train_x)
        train_y = torch.squeeze(train_y)
        self.train()
        likelihood.train()
        if optimizer == 'Adam':
            optimizer = torch.optim.Adam(self.parameters(), lr=0.1)
        else:
            logging.error('Undefined optimizer')
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, self)
        for i in range(training_iter):
            # Zero gradients from previous iteration
            optimizer.zero_grad()
            # Output from model
            output = self(train_x)
            # Calculate loss and backpropagate gradients
            loss = -mll(output, train_y)
            loss.backward()
            optimizer.step()

    def predict(self, x, likelihood, full_cov=False):
        self.eval()
        likelihood.eval()
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            observed_pred = likelihood(self(x))
            if not full_cov:
                mean = observed_pred.mean
                var = observed_pred.variance
            else:
                mean = observed_pred.mean
                var = observed_pred.covariance_matrix
        return mean, var

class BatchIndependentMultitaskGPModel(ExactGPModel):
    def __init__(self, train_x, train_y, likelihood, kernel, output_size=1):
        super().__init__(train_x, train_y, likelihood, kernel)
        self.mean_module = gpytorch.means.ConstantMean(
            batch_shape=torch.Size([output_size]))
        self.covar_module = kernel

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

if __name__ == '__main__':
    # Multi input, independent multi output
    train_x = torch.stack((
        torch.linspace(0, 1, 100), torch.linspace(0, 1, 100)), dim=1)
    train_y = torch.stack((
        torch.sin(train_x[:, 0] * (2 * math.pi)) + torch.randn(
            train_x[:, 0].size()) * math.sqrt(0.04),
        3 * torch.square(
            torch.cos(train_x[:, 1] * (2 * math.pi))) + torch.randn(
            train_x[:, 1].size()) * math.sqrt(0.0001),
        3 * torch.square(
            torch.cos(train_x[:, 1] * (2 * math.pi))) + torch.randn(
            train_x[:, 1].size()) * math.sqrt(0.0001)), dim=1)
    train_y = torch.squeeze(train_y)
    likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(
        num_tasks=train_y.shape[1],
        noise_prior=gpytorch.priors.MultivariateNormalPrior(
            torch.tensor([0.04, 0.04, 0.04]),
            torch.diag(torch.tensor([1., 1, 1]))))
    likelihood.noise_covar.noise = torch.tensor([0.04])
    likelihood.noise = torch.tensor([1e-4])
    lengthscale_prior = gpytorch.priors.MultivariateNormalPrior(
        torch.tensor([0.2, 0.2]),
        torch.diag(torch.tensor([6.0, 6.0])))
    outputscale_prior = gpytorch.priors.NormalPrior(1.0, 0.15)
    hypers = {'base_kernel.lengthscale': lengthscale_prior.mean,
              'outputscale': outputscale_prior.mean}
    kernel = gpytorch.kernels.ScaleKernel(
        gpytorch.kernels.RBFKernel(
            batch_shape=torch.Size([train_y.shape[1]]),
            ard_num_dims=train_x.shape[1]),
        batch_shape=torch.Size([train_y.shape[1]]))
    kernel.initialize(**hypers)  # init kernel hyperparams
    model = BatchIndependentMultitaskGPModel(train_x, train_y, likelihood,
                                             kernel,
                                             output_size=train_y.shape[1])
    print('multioutput')
    print(model.state_dict())
    # Evaluate GP by making predictions on test set
    test_x = torch.stack((
        torch.linspace(0, 1, 51), torch.linspace(0, 1, 51)), dim=1)
    mean, var = model.predict(test_x, likelihood)
    lower = mean - 2 * torch.sqrt(var)
    upper = mean + 2 * torch.sqrt(var)

** Stack trace/error message **

Traceback (most recent call last):
  File "/Users/mona/PhD_code/observer_add_GPyTorch_advanced/venv/lib/python3.9/site-packages/gpytorch/module.py", line 412, in __getattr__
    return super().__getattribute__(name)
AttributeError: 'MultitaskGaussianLikelihood' object has no attribute 'noise_covar'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/mona/PhD_code/observer_add_GPyTorch_advanced/src/script4.py", line 89, in <module>
    likelihood.noise_covar.noise = torch.tensor([0.04])
  File "/Users/mona/PhD_code/observer_add_GPyTorch_advanced/venv/lib/python3.9/site-packages/gpytorch/module.py", line 414, in __getattr__
    raise e
  File "/Users/mona/PhD_code/observer_add_GPyTorch_advanced/venv/lib/python3.9/site-packages/gpytorch/module.py", line 409, in __getattr__
    return super().__getattr__(name)
  File "/Users/mona/PhD_code/observer_add_GPyTorch_advanced/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 778, in __getattr__
    raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
torch.nn.modules.module.ModuleAttributeError: 'MultitaskGaussianLikelihood' object has no attribute 'noise_covar'

System information

• GPyTorch version 1.4.0
• PyTorch version 1.7.1
• Mac OS Big Sur

[Balandat]

just merged it

[wjmaddox]

Sorry for the confusion. I recently refactored MultitaskGaussianLikelihood to better exploit Kronecker structure throughout and probably should have done a better job documenting the changes. It looks like I need to add in a couple more setters for the task noise, which I’ll setup soon, as well as fixing the setter for the noise.

You should be able to set the raw task covariance noises by directly modifying the task noises likelihood.raw_task_noises.data = likelihood.raw_task_noises_constraint.inverse_transform(torch.tensor([0.04, 0.04, 0.04]))

Similarly, you can set likelihood.raw_noise.data = likelihood.raw_noise_constraint.inverse_transform(torch.tensor([1e-4]))