3D DGP Regression crashes from not enough RAM

[anja-sheppard]

Hi,

I'm trying to adapt the example here to work in a 2D grid space. However, after the model is trained and I try to predict with it, all 30 GB of RAM on my desktop is quickly eaten and then the code crashes. I don't think this should be such a RAM-intensive task, unless I misunderstand something. Any help is appreciated.

import os
import torch
import tqdm
import math
import gpytorch
from torch.nn import Linear
from gpytorch.means import ConstantMean, LinearMean
from gpytorch.kernels import MaternKernel, ScaleKernel
from gpytorch.variational import VariationalStrategy, CholeskyVariationalDistribution, LMCVariationalStrategy
from gpytorch.distributions import MultivariateNormal
from gpytorch.models.deep_gps import DeepGPLayer, DeepGP
from gpytorch.mlls import DeepApproximateMLL, VariationalELBO
from gpytorch.likelihoods import MultitaskGaussianLikelihood
from matplotlib import pyplot as plt

smoke_test = ('CI' in os.environ)
%matplotlib inline

train_x1 = torch.linspace(0, 1, 50)
train_x2 = torch.linspace(0, 1, 50)
train_x1_grid, train_x2_grid = torch.meshgrid(train_x1, train_x2)

train_Y = torch.stack([
    torch.sin(torch.sqrt(train_x1_grid**2 + train_x2_grid**2)),
    torch.cos(torch.sqrt(train_x1_grid**2 + train_x2_grid**2)),
    torch.sin((train_x1_grid * (2 * math.pi)) + torch.randn(train_x1_grid.size()) * 0.2) + ((train_x2_grid * (2 * math.pi)) + torch.randn(train_x2_grid.size()) * 0.2),
    torch.sin((train_x1_grid * (2 * math.pi)) + torch.randn(train_x1_grid.size()) * 0.2) + ((train_x2_grid * (2 * math.pi)) + torch.randn(train_x2_grid.size()) * 0.2)
], -1)

train_X = torch.stack([train_x1_grid.flatten(), train_x2_grid.flatten()]).T
train_Y = torch.flatten(train_Y, start_dim=0, end_dim=1)

class DGPHiddenLayer(DeepGPLayer):
    def __init__(self, input_dims, output_dims, num_inducing=128, linear_mean=True):
        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().__init__(variational_strategy, input_dims, output_dims)
        self.mean_module = ConstantMean() if linear_mean else LinearMean(input_dims)
        self.covar_module = ScaleKernel(
            MaternKernel(nu=2.5, 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)

num_tasks = train_Y.size(-1)
num_hidden_dgp_dims = 3


class MultitaskDeepGP(DeepGP):
    def __init__(self, train_x_shape):
        hidden_layer = DGPHiddenLayer(
            input_dims=train_x_shape[-1],
            output_dims=num_hidden_dgp_dims,
            linear_mean=True
        )
        last_layer = DGPHiddenLayer(
            input_dims=hidden_layer.output_dims,
            output_dims=num_tasks,
            linear_mean=False
        )

        super().__init__()

        self.hidden_layer = hidden_layer
        self.last_layer = last_layer

        # We're going to use a ultitask likelihood instead of the standard GaussianLikelihood
        self.likelihood = MultitaskGaussianLikelihood(num_tasks=num_tasks)

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

    def predict(self, test_x):
        with torch.no_grad():

            # The output of the model is a multitask MVN, where both the data points
            # and the tasks are jointly distributed
            # To compute the marginal predictive NLL of each data point,
            # we will call `to_data_independent_dist`,
            # which removes the data cross-covariance terms from the distribution.
            preds = model.likelihood(model(test_x)).to_data_independent_dist()

        return preds.mean.mean(0), preds.variance.mean(0)


model = MultitaskDeepGP(train_X.shape)

model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
mll = DeepApproximateMLL(VariationalELBO(model.likelihood, model, num_data=train_Y.size(0)))

num_epochs = 1 if smoke_test else 10
epochs_iter = tqdm.tqdm(range(num_epochs), desc="Epoch")
for i in epochs_iter:
    optimizer.zero_grad()
    output = model(train_X)
    loss = -mll(output, train_Y)
    epochs_iter.set_postfix(loss=loss.item())
    loss.backward()
    optimizer.step()

model.eval()

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    test_x1 = torch.linspace(0, 1, 10)
    test_x2 = torch.linspace(0, 1, 10)
    test_x1_grid, test_x2_grid = torch.meshgrid(test_x1, test_x2)
    test_X = torch.stack([train_x1_grid.flatten(), train_x2_grid.flatten()]).T
    mean, var = model.predict(test_X)  # CRASHES HERE