Improving the performance of gpytorch for the regression analysis

[neuronphysics]

Hi,
I should thank you first for this great library. My post is more asking help to achieve better results in my code. I used gpytorch for regression analysis or to be more accurate, I am trying to see whether I can use it to predict transition function for the robotic tasks in different the gym environments. The following code is my attempt to build this model. I used my model for the inverted pendulum

from torch.distributions import constraints
!pip install gpytorch
import tqdm
import gpytorch
from gpytorch.models.deep_gps import DeepGPLayer, DeepGP
from gpytorch.means import ConstantMean, ZeroMean, LinearMean
from gpytorch.kernels import RBFKernel, ScaleKernel, SpectralMixtureKernel
from gpytorch.variational import VariationalStrategy, CholeskyVariationalDistribution
from gpytorch.distributions import MultivariateNormal
from gpytorch.likelihoods import GaussianLikelihood, MultitaskGaussianLikelihood
from torch.utils.data import TensorDataset, DataLoader
from gpytorch.mlls import DeepApproximateMLL
from gpytorch.mlls import VariationalELBO
def hidden_size_extract(kwargs, name, delete_from_dict=False):
    if name not in kwargs:
        hidden_size = []
        for i in range(0, 6):
            key = name + '_%d' % i
            if key in kwargs and kwargs[key] != 0:
                hidden_size.append(kwargs[key])

                if delete_from_dict:
                    kwargs.pop(key)
    else:
        hidden_size = kwargs[name].copy()

        if delete_from_dict:
            kwargs.pop(name)

    return hidden_size

class DGPHiddenLayer(DeepGPLayer):

    def __init__(self, input_dims, output_dims, device, num_inducing, mean_type='constant', num_mixtures=5):
        if output_dims is None:
            inducing_points = torch.randn(num_inducing, input_dims).to(device=device)
            batch_shape = torch.Size([])
        else:
            inducing_points = torch.randn(output_dims, num_inducing, input_dims).to(device=device)
            batch_shape = torch.Size([output_dims])
        # Sparse Variational Formulation (inducing variables initialised as randn)
        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().__init__(variational_strategy, input_dims, output_dims)
        if mean_type == 'constant':
            self.mean_module = ConstantMean(batch_shape=batch_shape)
        elif mean_type == 'zero':
            self.mean_module = ZeroMean(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
        )


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

    def __call__(self, x, *other_inputs, **kwargs):

        if len(other_inputs):
            if isinstance(x, gpytorch.distributions.MultitaskMultivariateNormal):
                x = x.rsample()

            processed_inputs = [
                inp.unsqueeze(0).expand(gpytorch.settings.num_likelihood_samples.value(), *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)))



class DeepGaussianProcesses(DeepGP):
    def __init__(self, 
                 input_size, 
                 output_size, 
                 device, 
                 num_inducing=128,
                 noise_constraint=None, 
                 max_cholesky_size=10000,
                 optimizer = None,
                 **kwargs):
        super(DeepGaussianProcesses, self).__init__()
        # pass hidden layer sizes as separate arguments as well as array
        for k, v in kwargs.items():
            if k=='mean_type':
                setattr(self, k, v)
        hidden_size = hidden_size_extract(kwargs, 'hidden_size')
        
        self.output_size=output_size
        self.hidden_size = hidden_size
        self.hidden_size.append(output_size)
        self.max_cholesky_size = max_cholesky_size
        self.device = device
        if self.mean_type=='constant' and self.mean_type is not None:
           means = (len(self.hidden_size)-1)*['linear'] + ['constant'] # The last layer with constant mean
        else:
           means = (len(self.hidden_size)-1)*['linear'] + ['zero']
        hidden_layers = torch.nn.ModuleList([DGPHiddenLayer(
            input_dims=input_size,
            output_dims=self.hidden_size[0],
            device = self.device,
            mean_type=means[0],
            num_inducing=num_inducing,
            )])
        for i in range(len(self.hidden_size)-1):
            hidden_layers.append(DGPHiddenLayer(
                input_dims=self.hidden_size[i],
                output_dims=self.hidden_size[i + 1],
                device = self.device,
                mean_type=means[i+1],
                num_inducing = num_inducing,
                ))

        self.hidden_layers = hidden_layers
        # variable count of hidden layers and neurons

        #print(self._layers)
        if optimizer is not None:
           self._optimizer= optimizer
        else:
           self._optimizer= torch.optim.Adam([
               {'params': self.hidden_layers.parameters()}
            ], lr=1e-3, betas=(0.5,0.99))
           
        if output_size is None:
           self.likelihood = GaussianLikelihood(noise_constraint=noise_constraint)
        else:
           self.likelihood= MultitaskGaussianLikelihood(num_tasks=output_size, noise_constraint=noise_constraint)
        self.to(device=self.device)

    def forward(self, inputs):
        output = self.hidden_layers[0](inputs)
        for hid_layer in self.hidden_layers[1:]:
            output = hid_layer(output)
        return output

    def predict(self, test_loader):
        with torch.no_grad(), gpytorch.settings.max_cholesky_size(self.max_cholesky_size):
            mus = []
            variances = []
            samples = []
            lls = []
            lowers=[]
            uppers=[]
            for x_batch, y_batch in test_loader:
                preds = self.likelihood(self(x_batch))
                mus.append(preds.mean)
                variances.append(preds.variance)
                lls.append(self.likelihood.log_marginal(y_batch, self(x_batch)))
                lower, upper = preds.confidence_region()
                lowers.append(lower)
                uppers.append(upper)
                samples.append(preds.rsample())
        return torch.cat(mus, dim=-1), torch.cat(variances, dim=-1), torch.cat(lowers, dim=-1), torch.cat(uppers, dim=-1), torch.cat(samples, dim=-1), torch.cat(lls, dim=-1)


    def train_gp(self, train_x, train_y, num_epochs=20, num_samples=10, batch_size=200):
        scheduler = torch.optim.lr_scheduler.MultiStepLR(self._optimizer, milestones=[int(0.5 * num_epochs), int(0.75 * num_epochs)])
        train_dataset = TensorDataset(train_x, train_y)
        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        mll = DeepApproximateMLL(VariationalELBO(self.likelihood, self, num_data=len(train_x)))
        self.losses = []
        epochs_iter = tqdm.notebook.tqdm(range(num_epochs), desc="GP training Epoch")
        self.train()
        with gpytorch.settings.max_cholesky_size(self.max_cholesky_size):
             for i in epochs_iter:
                 # Within each iteration, we will go over each minibatch of data
                 minibatch_iter = tqdm.notebook.tqdm(train_loader, desc="Minibatch", leave=False)
            
                 for x_batch, y_batch in minibatch_iter:
                     with gpytorch.settings.num_likelihood_samples(num_samples):
                          self._optimizer.zero_grad()
                          output = self(x_batch)
                          #print(x_batch.shape, output, output.mean.shape, y_batch.shape)
                          loss = -mll(output, y_batch)
                          loss.backward()
                          self._optimizer.step()
                          scheduler.step()
                          self.losses.append(loss.item())
                          #print(self.hidden_layers[0].covar_module.outputscale)
                          minibatch_iter.set_postfix(loss=loss.item(),
                                                     ls1=self.hidden_layers[0].covar_module.base_kernel.lengthscale.mean().item(),
                                                     os1=self.hidden_layers[0].covar_module.outputscale.mean().item(),
                                                     ls2=self.hidden_layers[-1].covar_module.base_kernel.lengthscale.mean().item(),
                                                     os2=self.hidden_layers[-1].covar_module.outputscale.mean().item(),
                                                     noise=self.likelihood.noise.item(),
                                                    )
    def evaluation(self, train_x, train_y, test_x, test_y, plot=True, num_samples=10, batch_size=200):
        with gpytorch.settings.num_likelihood_samples(num_samples), torch.no_grad():
            self.eval()
            mll = gpytorch.mlls.VariationalELBO(self.likelihood, self, num_data=len(train_x))
            mll = gpytorch.mlls.DeepApproximateMLL(mll)
            nmll = 0.
            for x_batch, y_batch in zip(train_x.split(batch_size), train_y.split(batch_size)):
                output = self(x_batch)
                nmll += -mll(output, y_batch).mul(len(y_batch))
            nmll /= len(train_y)
        num_task = test_y.shape[1]
        with torch.no_grad():
            nll = 0.
            mse = 0.
            self.eval()
            mean_pred = torch.zeros(1, num_task)
            mean_true = torch.zeros(1, num_task)

            for x_batch, y_batch in zip(test_x.split(batch_size), test_y.split(batch_size)):
                pred_dist = self.likelihood(self(x_batch))
                mean_pred = torch.cat([mean_pred, pred_dist.mean.mean(0).cpu()])
                mean_true = torch.cat([mean_true, y_batch.cpu()])
                #torch.Size([800, 10]) torch.Size([4, 200, 10])
                pred_dist = torch.distributions.Normal(pred_dist.mean.T, pred_dist.stddev.T)
                mixture_weights = x_batch.new_ones(pred_dist.batch_shape[-1])
                pred_dist = torch.distributions.MixtureSameFamily(
                    mixture_distribution=torch.distributions.Categorical(mixture_weights),
                    component_distribution=pred_dist
                )
                nll += -pred_dist.log_prob(y_batch.T).sum(dim=-1).mean()
                mse += (pred_dist.mean - y_batch.T).pow(2).sum(dim=-1).mean()

            nll = nll.div(len(test_y))
            rmse = mse.div(len(test_y)).sqrt()
        mean_pred = mean_pred[1:, :]
        mean_true = mean_true[1:, :]
        print(nll, rmse,nmll)
        test_dataset = TensorDataset(test_x, test_y)
        test_loader = DataLoader(test_dataset, batch_size=200)
        predictive_means_test, predictive_variances_test, predictive_lower, predictive_upper, predictive_sample, test_lls_test = self.predict(test_loader)
       
        print(mean_pred.shape, test_y.shape)
      
        if plot:
          from matplotlib import pyplot as plt
          ncols=y.shape[-1]+1
          
          fig, axs = plt.subplots(nrows=1, ncols=ncols, figsize=(20, 7))

          axs = axs.ravel()

          for i in range(y.shape[-1]+1):
              if i==0:
                 axs[i].plot(np.array(self.losses))
                 axs[i].set_xlabel("iteration")
                 axs[i].set_ylabel("loss")
                 asp = np.diff(axs[i].get_xlim())[0] / np.diff(axs[i].get_ylim())[0]
                 axs[i].set_aspect(asp)
              else:               
                 axs[i].scatter(mean_true[:, i-1].numpy(), mean_pred[:, i-1].numpy())
                 axs[i].set_xlabel("true test state {}".format(i))
                 axs[i].set_ylabel("predict state {}".format(i))
                 asp = np.diff(axs[i].get_xlim())[0] / np.diff(axs[i].get_ylim())[0]
                 axs[i].set_aspect(asp)
                 #diag_line, = axs[i].plot(axs[i].get_xlim(), axs[i].get_ylim(), ls="--", c=".3")
                 diag_line, = axs[i].plot(axs[i].get_xlim(), axs[i].get_xlim(), ls="--", c=".3")
          fig.tight_layout()

        return nmll.item(), nll.item(), rmse.item()

if __name__=="__main__":
    from math import floor
    env=gym.make("InvertedPendulum-v2")
    state = env.reset()
    print(state.shape, env.action_space.shape[0], env.action_space.high)
    agent=Agent(input_dims=state.shape, env=env, n_actions=env.action_space.shape[0])
    n_games=600
    X = []
    y = []
    
    n = 10000
    for i in range(n):
        action=agent.choose_action(state)
        new_state, reward, done, _ = env.step(action)
        train_x_sample = torch.cat([torch.tensor(state).float(), torch.tensor(action).float()])
        X.append(train_x_sample)
        y.append(torch.tensor(new_state).float())
        state = new_state
        if done:
           state = env.reset()
    print("collectinng data is done .....")
    X = torch.stack(X, dim=0)
    print(X.shape)
    y = torch.stack(y, dim=0)
    print(y.shape)
    train_n = int(floor(0.8 * len(X)))
    train_x = X[:train_n, :].contiguous()
    train_y = y[:train_n].contiguous()
    test_x = X[train_n:, :].contiguous()
    test_y = y[train_n:].contiguous()
    if torch.cuda.is_available():
       train_x, train_y, test_x, test_y = train_x.cuda(), train_y.cuda(), test_x.cuda(), test_y.cuda()
    model = DeepGaussianProcesses(train_x.shape[-1], train_y.shape[-1], num_inducing=120, mean_type='constant', device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
    model.train_gp( train_x, train_y, num_epochs=350)
    nmll, nll, rmse=model.evaluation(train_x, train_y, test_x, test_y)

Here are the predictions versus true data. As you may notice at the end of some dimensions, the limits of the prediction is not exactly equal to the true data. They look like a sinusoidal shape tbh. I am wondering whether there is another covariance functions or mean, or even other small tweaks in the bits and pieces of the model which can help improve these results? Thanks in advance.

P.S. Is it possible to add some of the recurrent gaussian process algorithms to this library?

[gpleiss]

My suggestion would be to start with a simpler model (i.e. a non-deep GP) and see if this solves the errors. It's hard to pinpoint anything wrong that's in particular, but the model you have is a bit complicated.

[neuronphysics]

Actually, if you might noticed that I only use one layer. I didn't use many layers. I am wondering about another type of variational distributions such as TrilNaturalVariationalDistribution combined with LMCVariationalStrategy, since my data requires a multi output model? Do you think they would be better fit to the type of problem which I am tackling here?

[gpleiss]

I don't think the variational distribution is to blame, but it's worth a shot. Also maybe try more inducing points.