Repository: cornellius-gp/gpytorch
Pull request #668 (2019-04-26 17:36:54)

Implement Polynomial Kernel
Implements the polynomial kernel + unit tests

@dme65 


Patch:
diff --git a/docs/source/kernels.rst b/docs/source/kernels.rst
index f9acc0c0..dbe1e4a3 100644
--- a/docs/source/kernels.rst
+++ b/docs/source/kernels.rst
@@ -45,6 +45,12 @@ Standard Kernels
 .. autoclass:: PeriodicKernel
    :members:
 
+:hidden:`PolynomialKernel`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: PolynomialKernel
+   :members:
+
 :hidden:`RBFKernel`
 ~~~~~~~~~~~~~~~~~~~
 
diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
index 7ee10c48..254b8618 100644
--- a/gpytorch/kernels/__init__.py
+++ b/gpytorch/kernels/__init__.py
@@ -13,6 +13,7 @@ from .linear_kernel import LinearKernel
 from .matern_kernel import MaternKernel
 from .multitask_kernel import MultitaskKernel
 from .periodic_kernel import PeriodicKernel
+from .polynomial_kernel import PolynomialKernel
 from .product_structure_kernel import ProductStructureKernel
 from .rbf_kernel import RBFKernel
 from .rbf_kernel_grad import RBFKernelGrad
@@ -36,6 +37,7 @@ __all__ = [
     "MaternKernel",
     "MultitaskKernel",
     "PeriodicKernel",
+    "PolynomialKernel",
     "ProductKernel",
     "ProductStructureKernel",
     "RBFKernel",
diff --git a/gpytorch/kernels/polynomial_kernel.py b/gpytorch/kernels/polynomial_kernel.py
new file mode 100644
index 00000000..7cd2e3f8
--- /dev/null
+++ b/gpytorch/kernels/polynomial_kernel.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+
+import torch
+from .kernel import Kernel
+from typing import Optional
+from ..priors import Prior
+from ..constraints import Positive, Interval
+
+
+class PolynomialKernel(Kernel):
+    r"""
+    Computes a covariance matrix based on the Polynomial kernel
+    between inputs :math:`\mathbf{x_1}` and :math:`\mathbf{x_2}`:
+
+    .. math::
+        \begin{equation*}
+            k_\text{Poly}(\mathbf{x_1}, \mathbf{x_2}) = (\mathbf{x_1}^\top
+            \mathbf{x_2} + c)^{d}.
+        \end{equation*}
+
+    where
+
+    * :math:`c` is an :attr:`offset` parameter.
+
+    Args:
+        :attr:`offset_prior` (:class:`gpytorch.priors.Prior`):
+            Prior over the offset parameter (default `None`).
+        :attr:`offset_constraint` (Constraint, optional):
+            Constraint to place on offset parameter. Default: `Positive`.
+        :attr:`active_dims` (list):
+            List of data dimensions to operate on.
+            `len(active_dims)` should equal `num_dimensions`.
+    """
+
+    def __init__(
+        self,
+        power: int,
+        offset_prior: Optional[Prior] = None,
+        offset_constraint: Optional[Interval] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        if offset_constraint is None:
+            offset_constraint = Positive()
+
+        self.register_parameter(
+            name="raw_offset",
+            parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1))
+        )
+
+        # We want the power to be a float so we dont have to worry about its device / dtype.
+        if torch.is_tensor(power):
+            if power.numel() > 1:
+                raise RuntimeError("Cant create a Polynomial kernel with more than one power")
+            else:
+                power = power.item()
+
+        self.power = power
+
+        if offset_prior is not None:
+            self.register_prior(
+                "offset_prior",
+                offset_prior,
+                lambda: self.offset,
+                lambda v: self._set_offset(v)
+            )
+
+        self.register_constraint("raw_offset", offset_constraint)
+
+    @property
+    def offset(self) -> torch.Tensor:
+        return self.raw_offset_constraint.transform(self.raw_offset)
+
+    @offset.setter
+    def offset(self, value: torch.Tensor) -> None:
+        self._set_offset(value)
+
+    def _set_offset(self, value: torch.Tensor) -> None:
+        if not torch.is_tensor(value):
+            value = torch.as_tensor(value).to(self.raw_offset)
+        self.initialize(raw_offset=self.raw_offset_constraint.inverse_transform(value))
+
+    def forward(
+        self,
+        x1: torch.Tensor,
+        x2: torch.Tensor,
+        diag: Optional[bool] = False,
+        last_dim_is_batch: Optional[bool] = False,
+        **params
+    ) -> torch.Tensor:
+        offset = self.offset.view(*self.batch_shape, 1, 1)
+
+        if last_dim_is_batch:
+            x1 = x1.transpose(-1, -2).unsqueeze(-1)
+            x2 = x2.transpose(-1, -2).unsqueeze(-1)
+
+        if diag:
+            return ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)
+
+        if x1.dim() == 2 and x2.dim() == 2:
+            return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)
+        else:
+            return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)


Test patch:
diff --git a/test/kernels/test_polynomial_kernel.py b/test/kernels/test_polynomial_kernel.py
new file mode 100644
index 00000000..aab09fe4
--- /dev/null
+++ b/test/kernels/test_polynomial_kernel.py
@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+
+import torch
+import unittest
+from gpytorch.kernels import PolynomialKernel
+from test.kernels._base_kernel_test_case import BaseKernelTestCase
+
+
+class TestPolynomialKernel(unittest.TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return PolynomialKernel(power=2, **kwargs)
+
+    def test_computes_quadratic_kernel(self):
+        a = torch.tensor([[4, 1], [2, 2], [8, 0]], dtype=torch.float)
+        b = torch.tensor([[0, 0], [2, 1], [1, 0]], dtype=torch.float)
+        kernel = PolynomialKernel(power=2)
+        kernel.eval()
+
+        actual = torch.zeros(3, 3)
+        for i in range(3):
+            for j in range(3):
+                actual[i, j] = (a[i].matmul(b[j]) + kernel.offset).pow(kernel.power)
+
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = torch.zeros(2, 3, 3)
+        for l in range(2):
+            actual[l] = kernel(a[:, l].unsqueeze(-1), b[:, l].unsqueeze(-1)).evaluate()
+
+        res = kernel(a, b, last_dim_is_batch=True).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, b, last_dim_is_batch=True).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_cubic_kernel(self):
+        a = torch.tensor([[4, 1], [2, 2], [8, 0]], dtype=torch.float)
+        b = torch.tensor([[0, 0], [2, 1], [1, 0]], dtype=torch.float)
+        kernel = PolynomialKernel(power=3)
+        kernel.eval()
+
+        actual = torch.zeros(3, 3)
+        for i in range(3):
+            for j in range(3):
+                actual[i, j] = (a[i].matmul(b[j]) + kernel.offset).pow(kernel.power)
+
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = torch.zeros(2, 3, 3)
+        for l in range(2):
+            actual[l] = kernel(a[:, l].unsqueeze(-1), b[:, l].unsqueeze(-1)).evaluate()
+
+        res = kernel(a, b, last_dim_is_batch=True).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, b, last_dim_is_batch=True).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_quadratic_kernel_batch(self):
+        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
+        b = torch.tensor([[0, 2, 1], [-1, 2, 0]], dtype=torch.float).view(2, 3, 1)
+        kernel = PolynomialKernel(power=2, batch_shape=torch.Size([2])).initialize(offset=torch.rand(2, 1))
+        kernel.eval()
+
+        actual = torch.zeros(2, 3, 3)
+        for k in range(2):
+            for i in range(3):
+                for j in range(3):
+                    actual[k, i, j] = (a[k, i].matmul(b[k, j]) + kernel.offset[k]).pow(kernel.power)
+
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_cubic_kernel_batch(self):
+        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
+        b = torch.tensor([[0, 2, 1], [-1, 2, 0]], dtype=torch.float).view(2, 3, 1)
+        kernel = PolynomialKernel(power=3, batch_shape=torch.Size([2])).initialize(offset=torch.rand(2, 1))
+        kernel.eval()
+
+        actual = torch.zeros(2, 3, 3)
+        for k in range(2):
+            for i in range(3):
+                for j in range(3):
+                    actual[k, i, j] = (a[k, i].matmul(b[k, j]) + kernel.offset[k]).pow(kernel.power)
+
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+
+if __name__ == "__main__":
+    unittest.main()