Gpytorch: Fowarding through RBFKernel in batch-mode broken

[samuelstanton]

The recent changes to master have basically completely broken batch mode. This example is for the forward method in kernels.RBFKernel. As an example, consider the following code snippet

import torch
import gpytorch

class BatchGP(gpytorch.models.ExactGP):
    def __init__(self, train_inputs, train_targets, likelihood):
        super(BatchGP, self).__init__(train_inputs, train_targets, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_size=2)
        self.covar_module = gpytorch.kernels.RBFKernel(ard_num_dims=3, batch_size=2)
        self.register_parameter('log_outputscale', torch.nn.Parameter(torch.zeros(2, 1, 1)))

    def forward(self, inputs):
        inputs = inputs.unsqueeze(0)
        inputs = inputs.expand(2, -1, -1)
        mean = self.mean_module(inputs)
        prescale_covar = self.covar_module(inputs).evaluate()
        covar = self.log_outputscale.exp() * prescale_covar

        return self.likelihood(gpytorch.random_variables.GaussianRandomVariable(mean, covar))

train_x = torch.rand(5, 3)
train_y = torch.rand(5, 2)
likelihood = gpytorch.likelihoods.GaussianLikelihood()

batch_gp = BatchGP(train_x, train_y, likelihood)

batch_gp(train_x)

Out:

RuntimeError                              Traceback (most recent call last)
<ipython-input-11-bbf86dd62a62> in <module>()
     24 batch_gp = BatchGP(train_x, train_y, likelihood)
     25 
---> 26 batch_gp(train_x)

~/Code/gpytorch/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
     79                 if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
     80                     raise RuntimeError("You must train on the training inputs!")
---> 81             return super(ExactGP, self).__call__(*inputs, **kwargs)
     82 
     83         # Posterior mode

~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
    178 
    179     def __call__(self, *inputs, **kwargs):
--> 180         outputs = self.forward(*inputs, **kwargs)
    181         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyTensor):
    182             return outputs

<ipython-input-11-bbf86dd62a62> in forward(self, inputs)
     13         inputs = inputs.expand(2, -1, -1)
     14         mean = self.mean_module(inputs)
---> 15         prescale_covar = self.covar_module(inputs).evaluate()
     16         covar = self.log_outputscale.exp() * prescale_covar
     17 

~/Code/gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate(self)
    123 
    124     def evaluate(self):
--> 125         return self.evaluate_kernel().evaluate()
    126 
    127     def exact_predictive_mean(self, full_mean, train_labels, n_train, likelihood, precomputed_cache=None):

~/Code/gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
    104                 x2 = self.x2
    105 
--> 106             self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2, **self.params)
    107             if self.squeeze_row:
    108                 self._cached_kernel_eval.squeeze_(-2)

~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
    178 
    179     def __call__(self, *inputs, **kwargs):
--> 180         outputs = self.forward(*inputs, **kwargs)
    181         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyTensor):
    182             return outputs

~/Code/gpytorch/gpytorch/kernels/rbf_kernel.py in forward(self, x1, x2)
     94     def forward(self, x1, x2):
     95         x1_, x2_ = self._create_input_grid(x1, x2)
---> 96         x1_ = x1_.div(self.lengthscale)
     97         x2_ = x2_.div(self.lengthscale)
     98 

RuntimeError: The size of tensor a (5) must match the size of tensor b (2) at non-singleton dimension 1

If you go in and look at the dimension of the tensors, x1_.size() is [B x N x 1 x D] and self.lengthscale.size() is [B x 1 x D]

[Balandat]

@samuelstanton, #262 will use torch.distributions instead of the gpytorch RandomVariables - this should make sure that sample and log_prob work consistently in batch mode. Will update the PR tonight, hopefully we can get that in soon.

[gpleiss]

@samuelstanton - I think the issue is in your example? Your test_mean should be a 2 x 3 tensor.

[samuelstanton]

@gpleiss There was indeed a mistake in that example. I've update the issue with an example of the issue with forwarding through kernels in batch mode. After some more investigation it seems to be something that happens specifically in covar_module(input).evaluate(). If you don't try to convert to tensor the forward pass seems to work fine. e.g.

    def forward(self, inputs):
        inputs = inputs.unsqueeze(0)
        inputs = inputs.expand(2, -1, -1)
        mean = self.mean_module(inputs)
        covar = gpytorch.lazy.DiagLazyTensor(self.covar_module(inputs).diag())

        return self.likelihood(gpytorch.random_variables.GaussianRandomVariable(mean, covar))

But .evaluate() does work sometimes, like the following:

    def forward(self, inputs):
        inputs = inputs.unsqueeze(0)
        inputs = inputs.expand(2, -1, -1)
        mean = self.mean_module(inputs)
        prescaled_covar = gpytorch.lazy.DiagLazyTensor(self.covar_module(inputs).diag()).evaluate()
        covar = self.log_outputscale.exp() * prescaled_covar

        return self.likelihood(gpytorch.random_variables.GaussianRandomVariable(mean, covar))

So maybe a better name for the issue would be .evaluate() is unstable in batchmode?

[gpleiss]

@samuelstanton what's the exact error message you're getting?

[samuelstanton]

It's an error occurring in lazy_evaluated_kernel_tensor.py. The dimensions of the arguments to the kernel and the dimension of the lengthscale doesn't line up properly. If you go in and look at the dimension of the tensors, x1_.size() is [B x N x 1 x D] and self.lengthscale.size() is [B x 1 x D], so the .div call throws an error, as you can see in the example I gave.

[gpleiss]

Figured out the error. It was introduced by #266 😬
Fixing it right now

[gpleiss]

@samuelstanton try now

[samuelstanton]

That fixed that issue, but there appears to be another bug. If you try to run the following:

train_x = torch.rand(5, 3)
train_y = torch.rand(2, 5)
likelihood = gpytorch.likelihoods.GaussianLikelihood()
batch_gp = BatchGP(train_x, train_y, likelihood)
gp_output = batch_gp(train_x)
mll = gpytorch.mlls.ExactMarginalLogLikelihood(batch_gp.likelihood, batch_gp)
print(gp_output.mean().size())
print(train_y.size())
loss = -mll(gp_output, train_y)

You get

torch.Size([2, 5])
torch.Size([2, 5])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-3-d026ccc05135> in <module>()
     27 print(gp_output.mean().size())
     28 print(train_y.size())
---> 29 loss = -mll(gp_output, train_y)

~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
    178 
    179     def __call__(self, *inputs, **kwargs):
--> 180         outputs = self.forward(*inputs, **kwargs)
    181         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyTensor):
    182             return outputs

~/Code/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, output, target)
     49 
     50         # Get log determininat and first part of quadratic form
---> 51         inv_quad, log_det = covar.inv_quad_log_det(inv_quad_rhs=(target - mean).unsqueeze(-1), log_det=True)
     52 
     53         # Add terms for SGPR / when inducing points are learned

~/Code/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_log_det(self, inv_quad_rhs, log_det)
    593             inv_quad=(inv_quad_rhs is not None),
    594             log_det=log_det,
--> 595             preconditioner=self._preconditioner()[0],
    596             log_det_correction=self._preconditioner()[1],
    597         )(*args)

~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _preconditioner(self)
     44         if not hasattr(self, "_woodbury_cache"):
     45             max_iter = settings.max_preconditioner_size.value()
---> 46             self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_var, max_iter)
     47             self._woodbury_cache = pivoted_cholesky.woodbury_factor(self._piv_chol_self, self._diag_var.diag())
     48 

~/Code/gpytorch/gpytorch/utils/pivoted_cholesky.py in pivoted_cholesky(matrix, max_iter, error_tol)
     59                 row.unsqueeze_(0)
     60         else:
---> 61             row = matrix[full_batch_slice, pi_m, :]
     62 
     63         if isinstance(row, LazyTensor):

~/Code/gpytorch/gpytorch/lazy/constant_mul_lazy_tensor.py in __getitem__(self, i)
    148             first_index = i[0] if isinstance(i, tuple) else i
    149             constant = constant[first_index]
--> 150         return self.lazy_var.__getitem__(i) * constant

~/Code/gpytorch/gpytorch/lazy/lazy_tensor.py in __mul__(self, other)
   1024             return other
   1025 
-> 1026         return self.mul(other)
   1027 
   1028     def __getitem__(self, index):

~/Code/gpytorch/gpytorch/lazy/lazy_tensor.py in mul(self, other)
    669                 "Expected: size of [1] or [%d] or %s.\n"
    670                 "Got: size of %s"
--> 671                 % (self.size(0) if self.ndimension() == 3 else 1, repr(self.size()), repr(other.size()))
    672             )
    673 

RuntimeError: "other" must be a constant (or batch of constants), or the same size as self.
Expected: size of [1] or [1] or torch.Size([2, 5]).
Got: size of torch.Size([2])

What's weird about this one, is it seems like you shouldn't be getting it if train_targets are the same size as output.mean() which you can see that they are. Should I open another issue for this? On another note, I'm going to convert this script to a unit-test.

[gpleiss]

You should use ScaleKernel(RBFKernel()) for outputscales - I think then it should work.