Repository: cornellius-gp/gpytorch
Pull request #345 (2018-11-04 01:54:26)

Improve GridKernel
GridKernel and GridInterpolationKernel both constructed the full dimensional grid of points despite clearly not needing them to function. I think this was an artifact from back when variational inference needed to know these things?

Either way, we obviously shouldn't be doing this because it requires O(dm^d) storage to store `inducing_points` compared to O(dm) storage to store `grid`.

Related to #344, this also dramatically simplifies usage of `GridKernel` for data that exactly lies on a grid by requiring a user to only specify the grid in each dimension.


Patch:
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
index 806c38df..32e96085 100644
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -91,10 +91,10 @@ class GridInterpolationKernel(GridKernel):
         self.num_dims = num_dims
         self.grid_size = grid_size
         self.grid_bounds = grid_bounds
-        inducing_points, grid = self._create_grid()
+        grid = self._create_grid()
 
         super(GridInterpolationKernel, self).__init__(
-            base_kernel=base_kernel, inducing_points=inducing_points, grid=grid, active_dims=active_dims
+            base_kernel=base_kernel, grid=grid, active_dims=active_dims
         )
         self.register_buffer("has_initialized_grid", torch.tensor(has_initialized_grid, dtype=torch.uint8))
 
@@ -106,16 +106,7 @@ class GridInterpolationKernel(GridKernel):
                 self.grid_bounds[i][0] - grid_diff, self.grid_bounds[i][1] + grid_diff, self.grid_size
             )
 
-        inducing_points = torch.zeros(int(pow(self.grid_size, len(self.grid_bounds))), len(self.grid_bounds))
-        prev_points = None
-        for i in range(len(self.grid_bounds)):
-            for j in range(self.grid_size):
-                inducing_points[j * self.grid_size ** i : (j + 1) * self.grid_size ** i, i].fill_(grid[j, i])
-                if prev_points is not None:
-                    inducing_points[j * self.grid_size ** i : (j + 1) * self.grid_size ** i, :i].copy_(prev_points)
-            prev_points = inducing_points[: self.grid_size ** (i + 1), : (i + 1)]
-
-        return inducing_points, grid
+        return grid
 
     @property
     def _tight_grid_bounds(self):
@@ -145,7 +136,7 @@ class GridInterpolationKernel(GridKernel):
 
     def _inducing_forward(self, batch_dims, **params):
         return super(GridInterpolationKernel, self).forward(
-            self.inducing_points, self.inducing_points, batch_dims=batch_dims, **params
+            self.grid, self.grid, batch_dims=batch_dims, **params
         )
 
     def forward(self, x1, x2, batch_dims=None, **params):
@@ -173,8 +164,8 @@ class GridInterpolationKernel(GridKernel):
                     (x_min - 2.01 * spacing, x_max + 2.01 * spacing)
                     for x_min, x_max, spacing in zip(x_mins, x_maxs, grid_spacings)
                 )
-                inducing_points, grid = self._create_grid()
-                self.update_inducing_points_and_grid(inducing_points, grid)
+                grid = self._create_grid()
+                self.update_grid(grid)
 
         base_lazy_tsr = self._inducing_forward(batch_dims=batch_dims, **params)
         if x1.size(0) > 1:
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index 0302b5c9..e6c8dacc 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -25,10 +25,6 @@ class GridKernel(Kernel):
     Args:
         :attr:`base_kernel` (Kernel):
             The kernel to speed up with grid methods.
-        :attr:`inducing_points` (Tensor, n x d):
-            This will be the set of points that lie on the grid.
-        :attr:`grid` (Tensor, k x d):
-            The exact grid points.
         :attr:`active_dims` (tuple of ints, optional):
             Passed down to the `base_kernel`.
 
@@ -36,12 +32,9 @@ class GridKernel(Kernel):
         http://www.cs.cmu.edu/~andrewgw/manet.pdf
     """
 
-    def __init__(self, base_kernel, inducing_points, grid, active_dims=None):
+    def __init__(self, base_kernel, grid, active_dims=None):
         super(GridKernel, self).__init__(active_dims=active_dims)
         self.base_kernel = base_kernel
-        if inducing_points.ndimension() != 2:
-            raise RuntimeError("Inducing points should be 2 dimensional")
-        self.register_buffer("inducing_points", inducing_points.unsqueeze(0))
         self.register_buffer("grid", grid)
 
     def train(self, mode=True):
@@ -49,26 +42,39 @@ class GridKernel(Kernel):
             del self._cached_kernel_mat
         return super(GridKernel, self).train(mode)
 
-    def update_inducing_points_and_grid(self, inducing_points, grid):
+    def update_grid(self, grid):
         """
-        Supply a new set of `inducing_points` and a new `grid` if they ever change.
+        Supply a new `grid` if it ever changes.
         """
-        self.inducing_points.detach().resize_(inducing_points.size()).copy_(inducing_points)
         self.grid.detach().resize_(grid.size()).copy_(grid)
         if hasattr(self, "_cached_kernel_mat"):
             del self._cached_kernel_mat
         return self
 
+    def _create_data_from_grid(self):
+        grid_size = self.grid.size(-2)
+        grid_dim = self.grid.size(-1)
+        grid_data = torch.zeros(int(pow(grid_size, grid_dim)), grid_dim)
+        prev_points = None
+        for i in range(grid_dim):
+            for j in range(grid_size):
+                grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(self.grid[j, i])
+                if prev_points is not None:
+                    grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, :i].copy_(prev_points)
+            prev_points = grid_data[: grid_size ** (i + 1), : (i + 1)]
+
+        return grid_data
+
     def forward(self, x1, x2, diag=False, batch_dims=None, **params):
-        if not torch.equal(x1, self.inducing_points) or not torch.equal(x2, self.inducing_points):
-            raise RuntimeError("The kernel should only receive the inducing points as input")
+        grid = self.grid.unsqueeze(0)
+        x1 = x1.unsqueeze(0) if x1.dim() == 2 else x1
+        x2 = x2.unsqueeze(0) if x2.dim() == 2 else x2
+        if torch.equal(x1, grid) and torch.equal(x2, grid):
 
-        if not self.training and hasattr(self, "_cached_kernel_mat"):
-            covar = self._cached_kernel_mat
+            if not self.training and hasattr(self, "_cached_kernel_mat"):
+                return self._cached_kernel_mat
 
-        else:
             n_dim = x1.size(-1)
-            grid = self.grid.unsqueeze(0)
 
             if settings.use_toeplitz.on():
                 first_item = grid[:, 0:1]
@@ -93,4 +99,8 @@ class GridKernel(Kernel):
             if not self.training:
                 self._cached_kernel_mat = covar
 
-        return covar
+            return covar
+        else:
+            x1_ = self._create_data_from_grid() if torch.equal(x1, grid) else x1
+            x2_ = self._create_data_from_grid() if torch.equal(x2, grid) else x2
+            return self.base_kernel.forward(x1_, x2_, diag=diag, batch_dims=batch_dims, **params)
diff --git a/gpytorch/utils/interpolation.py b/gpytorch/utils/interpolation.py
index 47f6d60d..f226f8c4 100644
--- a/gpytorch/utils/interpolation.py
+++ b/gpytorch/utils/interpolation.py
@@ -100,9 +100,6 @@ class Interpolation(object):
             lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()
             lower_grid_pt_idxs.detach_()
 
-            if len(lower_grid_pt_idxs.shape) == 0:
-                lower_grid_pt_idxs = lower_grid_pt_idxs.unsqueeze(0)
-
             scaled_dist = lower_pt_rel_dists.unsqueeze(-1) + interp_points_flip.unsqueeze(-2)
             dim_interp_values = self._cubic_interpolation_kernel(scaled_dist)
 


Test patch:
diff --git a/test/kernels/test_grid_kernel.py b/test/kernels/test_grid_kernel.py
new file mode 100644
index 00000000..cea36b62
--- /dev/null
+++ b/test/kernels/test_grid_kernel.py
@@ -0,0 +1,51 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import torch
+import unittest
+from gpytorch.kernels import RBFKernel, GridKernel, GridInterpolationKernel
+
+cv = GridInterpolationKernel(RBFKernel(), grid_size=10, grid_bounds=[(0, 1), (0, 2)])
+grid = cv.grid
+
+grid_size = grid.size(-2)
+grid_dim = grid.size(-1)
+grid_data = torch.zeros(int(pow(grid_size, grid_dim)), grid_dim)
+prev_points = None
+for i in range(grid_dim):
+    for j in range(grid_size):
+        grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(grid[j, i])
+        if prev_points is not None:
+            grid_data[j * grid_size ** i : (j + 1) * grid_size ** i, :i].copy_(prev_points)
+    prev_points = grid_data[: grid_size ** (i + 1), : (i + 1)]
+
+
+class TestGridKernel(unittest.TestCase):
+    def test_grid_grid(self):
+        base_kernel = RBFKernel()
+        kernel = GridKernel(base_kernel, grid)
+        grid_eval = kernel(grid, grid).evaluate()
+        actual_eval = base_kernel(grid_data, grid_data).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 1e-5)
+
+    def test_nongrid_grid(self):
+        base_kernel = RBFKernel()
+        data = torch.randn(5, 2)
+        kernel = GridKernel(base_kernel, grid)
+        grid_eval = kernel(grid, data).evaluate()
+        actual_eval = base_kernel(grid_data, data).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 1e-5)
+
+    def test_nongrid_nongrid(self):
+        base_kernel = RBFKernel()
+        data = torch.randn(5, 2)
+        kernel = GridKernel(base_kernel, grid)
+        grid_eval = kernel(data, data).evaluate()
+        actual_eval = base_kernel(data, data).evaluate()
+        self.assertLess(torch.norm(grid_eval - actual_eval), 1e-5)
+
+
+if __name__ == "__main__":
+    unittest.main()