"Memory issue in Multitask Gaussian Process Inference for Multioutput Problems with high dimensional data"

[saavnkr]

Hi! I'm new to Gpytorch switching from other packages because of its seamless integration with pytorch.

I'm trying to work on a Multioutput problem (using MultitaskGP, Exact Inference).
My data's input and output shape is (N,D), where N is the number of different samples with D being flattened pixels.
After successfully training my model when I want to make predictions or print the shape of either covariance or variance matrix
it throws out of memory error, and the memory requirement is way too high, to be specific, 100+ GB on GPU.
Instead of evaluating the full covariance matrix, is there any option to evaluate the (N,D) shaped matrix, instead of (N,D,N,D) which will be computationally less expensive?

[Balandat]

Can you share the code you're using? Note that in eval() mode, if you're passing in a n_test x d tensor, the model will compute the joint covariance over the n_test points. This is usually done lazily though, so as long as you only access the diagonal (i.e. marginal variances) this should be ok, unless your n_test is massive. You can also just compute the marginal distribution for each data point by passing in a batched tensor of shape n_test x 1 x d.

[saavnkr]

'n_test x d' is of shape (200,64*64). My goal is to find the standard deviation for each of the predictions for which I need the covariance matrix.

`
class GaussianProcessRegression(gpytorch.models.ExactGP):
def init(self, train_x, train_y, likelihood):
super().init(train_x, train_y, likelihood)
num_dims = train_y.shape[-1] # Number of output dimensions

    self.mean_module = gpytorch.means.MultitaskMean(
        gpytorch.means.ConstantMean(), num_tasks=num_dims
    )
            
    self.covar_module = gpytorch.kernels.MultitaskKernel(
        gpytorch.kernels.MaternKernel(), num_tasks=num_dims, rank=1, diag=True
    )

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

`
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=y_tr.shape[-1])
model = GaussianProcessRegression(x_tr, y_tr, likelihood)

model=model.to(device)
likelihood=likelihood.to(device)

model.train()
likelihood.train()

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood,model)

num_iterations = 80
for i in range(num_iterations):
optimizer.zero_grad()
output = model(x_tr)
loss = -mll(output, y_tr)
loss.backward()
print('Iter %d/%d - Loss: %.3f' % (i + 1, num_iterations, loss.item()))
optimizer.step()

model.eval()
likelihood.eval()

with torch.no_grad():
observed_pred = likelihood(model(x_t))

mean_pred = observed_pred.mean
var_pred = observed_pred.variance
`

This is the error message I get:

---

RuntimeError Traceback (most recent call last)
/tmp/ipykernel_3952697/2191948531.py in
1 # Make predictions for the test points
2 with torch.no_grad():
----> 3 observed_pred = likelihood(model(x_t))
4
5 # Access mean and covariance of the predicted distribution

.../site-packages/gpytorch/models/exact_gp.py) in call(self, *args, **kwargs)
330 predictive_mean,
331 predictive_covar,
--> 332 ) = self.prediction_strategy.exact_prediction(full_mean, full_covar)
333
334 # Reshape predictive mean to match the appropriate event shape

.../site-packages/gpytorch/models/exact_prediction_strategies.py) in exact_prediction(self, joint_mean, joint_covar)
271 return (
272 self.exact_predictive_mean(test_mean, test_train_covar),
--> 273 self.exact_predictive_covar(test_test_covar, test_train_covar),
274 )
275

.../site-packages/gpytorch/models/exact_prediction_strategies.py) in exact_predictive_covar(self, test_test_covar, test_train_covar)
317 train_train_covar = self.likelihood(dist, self.train_inputs).lazy_covariance_matrix
...
-> 2485 eye = torch.eye(num_cols, dtype=self.dtype, device=self.device)
2486 eye = eye.expand(*self.batch_shape, num_cols, num_cols)
2487 res = self.matmul(eye)

RuntimeError: CUDA out of memory. Tried to allocate 2500.00 GiB