Repository: cornellius-gp/gpytorch
Pull request #928 (2019-11-08 17:32:38)

Performance improvements in small data regimes (n < 2000)
In the past few months a lot of our speed optimizations for exact GPs have drifted towards favoring large data regimes (n > 100000). This PR makes two optimizations that achieve around a 2.5x speed up in lower data regimes when using CG for training:
1. Preconditioning is turned off until n > 2000.
1. Use torch.symeig instead of our own batch_symeig now that it supports batching.

This pushes the value of n for which CG overtakes Cholesky in terms of speed back down to around n=1000.


Patch:
diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
index 1e75bd35..49903d7e 100644
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -51,7 +51,7 @@ class AddedDiagLazyTensor(SumLazyTensor):
             return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)
 
     def _preconditioner(self):
-        if settings.max_preconditioner_size.value() == 0:
+        if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():
             return None, None, None
 
         if not (hasattr(self, "_woodbury_cache") or hasattr(self, "self._q_cache")):
diff --git a/gpytorch/settings.py b/gpytorch/settings.py
index 86fd940b..7c4290dc 100644
--- a/gpytorch/settings.py
+++ b/gpytorch/settings.py
@@ -404,6 +404,17 @@ class memory_efficient(_feature_flag):
     _state = False
 
 
+class min_preconditioning_size(_value_context):
+    """
+    If the size of of a LazyTensor is less than `min_preconditioning_size`,
+    then we won't use pivoted Cholesky based preconditioning.
+
+    Default: 2000
+    """
+
+    _global_value = 2000
+
+
 class num_likelihood_samples(_value_context):
     """
     The number of samples to draw from a latent GP when computing a likelihood
diff --git a/gpytorch/utils/__init__.py b/gpytorch/utils/__init__.py
index 1535f539..9a2dd7a7 100644
--- a/gpytorch/utils/__init__.py
+++ b/gpytorch/utils/__init__.py
@@ -5,7 +5,6 @@ from .linear_cg import linear_cg
 from .stochastic_lq import StochasticLQ
 from . import broadcasting
 from . import cholesky
-from . import eig
 from . import fft
 from . import grid
 from . import interpolation
@@ -33,7 +32,6 @@ __all__ = [
     "linear_cg",
     "StochasticLQ",
     "cholesky",
-    "eig",
     "fft",
     "grid",
     "interpolation",
diff --git a/gpytorch/utils/eig.py b/gpytorch/utils/eig.py
deleted file mode 100644
index d3e45313..00000000
--- a/gpytorch/utils/eig.py
+++ /dev/null
@@ -1,30 +0,0 @@
-#!/usr/bin/env python3
-
-import torch
-
-
-def batch_symeig(mat):
-    """
-    """
-    mat_orig = mat
-    dtkwargs = {"device": mat.device, "dtype": mat.dtype}
-    batch_shape = mat_orig.shape[:-2]
-    matrix_shape = mat_orig.shape[-2:]
-
-    # Smaller matrices are faster on the CPU than the GPU
-    if mat.size(-1) <= 32:
-        mat = mat.cpu()
-
-    mat = mat.view(-1, *matrix_shape)
-    eigenvectors = torch.empty(batch_shape.numel(), *matrix_shape, dtype=mat.dtype, device=mat.device)
-    eigenvalues = torch.empty(batch_shape.numel(), matrix_shape[-1], dtype=mat.dtype, device=mat.device)
-
-    for i in range(batch_shape.numel()):
-        evals, evecs = mat[i].symeig(eigenvectors=True)
-        mask = evals.ge(0)
-        eigenvectors[i] = evecs * mask.type_as(evecs).unsqueeze(0)
-        eigenvalues[i] = evals.masked_fill_(~mask, 1)
-
-    eigenvalues = eigenvalues.to(**dtkwargs).view(*batch_shape, -1)
-    eigenvectors = eigenvectors.to(**dtkwargs).view_as(mat_orig)
-    return eigenvalues, eigenvectors
diff --git a/gpytorch/utils/lanczos.py b/gpytorch/utils/lanczos.py
index 474d5eff..52aadf9e 100644
--- a/gpytorch/utils/lanczos.py
+++ b/gpytorch/utils/lanczos.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import torch
-from .eig import batch_symeig
 from .. import settings
 
 
@@ -160,4 +159,15 @@ def lanczos_tridiag_to_diag(t_mat):
 
     TODO: make the eigenvalue computations done in batch mode.
     """
-    return batch_symeig(t_mat)
+    orig_device = t_mat.device
+    if t_mat.size(-1) < 32:
+        retr = torch.symeig(t_mat.cpu(), eigenvectors=True)
+    else:
+        retr = torch.symeig(t_mat, eigenvectors=True)
+
+    evals, evecs = retr
+    mask = evals.ge(0)
+    evecs = evecs * mask.type_as(evecs).unsqueeze(-2)
+    evals = evals.masked_fill_(~mask, 1)
+
+    return evals.to(orig_device), evecs.to(orig_device)
diff --git a/gpytorch/utils/linear_cg.py b/gpytorch/utils/linear_cg.py
index 9b3f71b7..5e425eaf 100644
--- a/gpytorch/utils/linear_cg.py
+++ b/gpytorch/utils/linear_cg.py
@@ -159,7 +159,7 @@ def linear_cg(
 
     # Check for NaNs
     if not torch.equal(residual, residual):
-        raise RuntimeError("NaNs encounterd when trying to perform matrix-vector multiplication")
+        raise RuntimeError("NaNs encountered when trying to perform matrix-vector multiplication")
 
     # Sometime we're lucky and the preconditioner solves the system right away
     # Check for convergence


Test patch:
diff --git a/test/examples/test_kissgp_multiplicative_regression.py b/test/examples/test_kissgp_multiplicative_regression.py
index 56d62b8a..fb829fef 100644
--- a/test/examples/test_kissgp_multiplicative_regression.py
+++ b/test/examples/test_kissgp_multiplicative_regression.py
@@ -14,6 +14,7 @@ from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
 from gpytorch.distributions import MultivariateNormal
+from gpytorch import settings
 
 # Simple training data: let's try to learn a sine function,
 # but with KISS-GP let's use 100 training examples.
@@ -83,13 +84,13 @@ class TestKISSGPMultiplicativeRegression(unittest.TestCase):
             loss.backward()
             optimizer.n_iter += 1
 
-        for param in gp_model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-        for param in likelihood.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-        optimizer.step()
+            for param in gp_model.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+            for param in likelihood.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+            optimizer.step()
 
         # Test the model
         gp_model.eval()
diff --git a/test/utils/test_pivoted_cholesky.py b/test/utils/test_pivoted_cholesky.py
index 3b56604f..da22d857 100644
--- a/test/utils/test_pivoted_cholesky.py
+++ b/test/utils/test_pivoted_cholesky.py
@@ -9,6 +9,7 @@ import torch
 from gpytorch.kernels import RBFKernel
 from gpytorch.test.utils import approx_equal
 from gpytorch.utils import pivoted_cholesky, woodbury
+from gpytorch import settings
 
 
 class TestPivotedCholesky(unittest.TestCase):
@@ -54,13 +55,13 @@ class TestPivotedCholesky(unittest.TestCase):
         size = 50
         X = torch.rand((size, 2)).to(dtype=dtype)
         y = torch.sin(torch.sum(X, 1)).unsqueeze(-1).to(dtype=dtype)
+        with settings.min_preconditioning_size(0):
+            noise = torch.DoubleTensor(size).uniform_(math.log(1e-3), math.log(1e-1)).exp_().to(dtype=dtype)
+            lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
+            precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()
 
-        noise = torch.DoubleTensor(size).uniform_(math.log(1e-3), math.log(1e-1)).exp_().to(dtype=dtype)
-        lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
-        precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()
-
-        F = lazy_tsr._piv_chol_self
-        M = noise.diag() + F.matmul(F.t())
+            F = lazy_tsr._piv_chol_self
+            M = noise.diag() + F.matmul(F.t())
 
         x_exact = torch.solve(y, M)[0]
         x_qr = precondition_qr(y)
@@ -75,11 +76,12 @@ class TestPivotedCholesky(unittest.TestCase):
         X = torch.rand((size, 2)).to(dtype=dtype)
         y = torch.sin(torch.sum(X, 1)).unsqueeze(-1).to(dtype=dtype)
 
-        noise = 1e-2 * torch.ones(size, dtype=dtype)
-        lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
-        precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()
+        with settings.min_preconditioning_size(0):
+            noise = 1e-2 * torch.ones(size, dtype=dtype)
+            lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
+            precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()
 
-        F = lazy_tsr._piv_chol_self
+            F = lazy_tsr._piv_chol_self
         M = noise.diag() + F.matmul(F.t())
 
         x_exact = torch.solve(y, M)[0]