Merge pull request #1709 from cornellius-gp/fix_sgpr_no_fast_comp

Fix SGPR prediction speed with no fast_computations

Author: jacobrgardner

gpytorch/lazy/low_rank_root_added_diag_lazy_tensor.py

@@ -124,3 +124,32 @@ def inv_quad_logdet(self, inv_quad_rhs=None, logdet=False, reduce_inv_quad=True)
             logdet_term = self._logdet()
 
         return inv_quad_term, logdet_term
+
+    def inv_matmul(self, right_tensor, left_tensor=None):
+        if not self.is_square:
+            raise RuntimeError(
+                "inv_matmul only operates on (batches of) square (positive semi-definite) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
+        if self.dim() == 2 and right_tensor.dim() == 1:
+            if self.shape[-1] != right_tensor.numel():
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, right_tensor.shape
+                    )
+                )
+
+        squeeze_solve = False
+        if right_tensor.ndimension() == 1:
+            right_tensor = right_tensor.unsqueeze(-1)
+            squeeze_solve = True
+
+        solve = self._solve(right_tensor)
+        if squeeze_solve:
+            solve = solve.squeeze(-1)
+
+        if left_tensor is not None:
+            return left_tensor @ solve
+        else:
+            return solve

test/examples/test_sgpr_regression.py

@@ -3,6 +3,7 @@
 import os
 import random
 import unittest
+from unittest.mock import MagicMock, patch
 import warnings
 from math import exp, pi
 
@@ -14,6 +15,7 @@
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
 from gpytorch.test.utils import least_used_cuda_device
+from gpytorch.utils.cholesky import CHOLESKY_METHOD
 from gpytorch.utils.warnings import NumericalWarning
 from torch import optim
 
@@ -65,41 +67,56 @@ def tearDown(self):
         if hasattr(self, "rng_state"):
             torch.set_rng_state(self.rng_state)
 
-    def test_sgpr_mean_abs_error(self):
+    def test_sgpr_mean_abs_error(self, cuda=False):
         # Suppress numerical warnings
         warnings.simplefilter("ignore", NumericalWarning)
 
-        train_x, train_y, test_x, test_y = make_data()
+        train_x, train_y, test_x, test_y = make_data(cuda=cuda)
         likelihood = GaussianLikelihood()
         gp_model = GPRegressionModel(train_x, train_y, likelihood)
         mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
 
-        # Optimize the model
-        gp_model.train()
-        likelihood.train()
+        if cuda:
+            gp_model = gp_model.cuda()
+            likelihood = likelihood.cuda()
 
-        optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
-        for _ in range(30):
-            optimizer.zero_grad()
-            output = gp_model(train_x)
-            loss = -mll(output, train_y)
-            loss.backward()
-            optimizer.step()
+        # Mock cholesky
+        _wrapped_cholesky = MagicMock(
+            wraps=torch.linalg.cholesky if CHOLESKY_METHOD == "torch.linalg.cholesky" else torch.linalg.cholesky_ex
+        )
+        with patch(CHOLESKY_METHOD, new=_wrapped_cholesky) as cholesky_mock:
 
-            # Check that we have the right LazyTensor type
-            kernel = likelihood(gp_model(train_x)).lazy_covariance_matrix.evaluate_kernel()
-            self.assertIsInstance(kernel, gpytorch.lazy.LowRankRootAddedDiagLazyTensor)
+            # Optimize the model
+            gp_model.train()
+            likelihood.train()
 
-        for param in gp_model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
+            optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
+            for _ in range(30):
+                optimizer.zero_grad()
+                output = gp_model(train_x)
+                loss = -mll(output, train_y)
+                loss.backward()
+                optimizer.step()
 
-        # Test the model
-        gp_model.eval()
-        likelihood.eval()
+                # Check that we have the right LazyTensor type
+                kernel = likelihood(gp_model(train_x)).lazy_covariance_matrix.evaluate_kernel()
+                self.assertIsInstance(kernel, gpytorch.lazy.LowRankRootAddedDiagLazyTensor)
 
-        test_preds = likelihood(gp_model(test_x)).mean
-        mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
+            for param in gp_model.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+
+            # Test the model
+            gp_model.eval()
+            likelihood.eval()
+
+            test_preds = likelihood(gp_model(test_x)).mean
+            mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
+            cholesky_mock.assert_called()  # We SHOULD call Cholesky...
+            for chol_arg in cholesky_mock.call_args_list:
+                first_arg = chol_arg[0][0]
+                self.assertTrue(torch.is_tensor(first_arg))
+                self.assertTrue(first_arg.size(-1) == gp_model.covar_module.inducing_points.size(-2))
 
         self.assertLess(mean_abs_error.squeeze().item(), 0.1)
 
@@ -123,62 +140,9 @@ def test_sgpr_mean_abs_error_cuda(self):
 
         if not torch.cuda.is_available():
             return
-        with least_used_cuda_device():
-            train_x, train_y, test_x, test_y = make_data(cuda=True)
-            likelihood = GaussianLikelihood().cuda()
-            gp_model = GPRegressionModel(train_x, train_y, likelihood).cuda()
-            mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
-
-            # Test the model before optimization
-            gp_model.eval()
-            likelihood.eval()
-            test_preds = likelihood(gp_model(test_x)).mean
-            mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
-            self.assertLess(mean_abs_error.squeeze().item(), 0.02)
-
-            # Test variances before optimization
-            test_vars = likelihood(gp_model(test_x)).variance
-            self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-            self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
-            self.assertLess(
-                test_vars.max().item() - 0.05,
-                likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
-            )
-
-            # Optimize the model
-            gp_model.train()
-            likelihood.train()
-
-            optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
-            optimizer.n_iter = 0
-            for _ in range(25):
-                optimizer.zero_grad()
-                output = gp_model(train_x)
-                loss = -mll(output, train_y)
-                loss.backward()
-                optimizer.n_iter += 1
-                optimizer.step()
 
-            for param in gp_model.parameters():
-                self.assertTrue(param.grad is not None)
-                self.assertGreater(param.grad.norm().item(), 0)
-
-            # Test the model
-            gp_model.eval()
-            likelihood.eval()
-            test_preds = likelihood(gp_model(test_x)).mean
-            mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
-
-            self.assertLess(mean_abs_error.squeeze().item(), 0.02)
-
-            # Test variances
-            test_vars = likelihood(gp_model(test_x)).variance
-            self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-            self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
-            self.assertLess(
-                test_vars.max().item() - 0.05,
-                likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
-            )
+        with least_used_cuda_device():
+            self.test_sgpr_mean_abs_error(cuda=True)
 
 
 if __name__ == "__main__":