Gaussian Process tutorial

.gitignore

@@ -1,3 +1,4 @@
+# temp
 run_outputs*
 .DS_Store
 data

docs/source/contrib.gp.rst

@@ -0,0 +1,5 @@
+.. automodule:: pyro.contrib.gp
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource

docs/source/contrib.rst

@@ -8,3 +8,4 @@ Contributed Code
    :maxdepth: 2
 
    contrib.named
+   contrib.gp

pyro/contrib/__init__.py

@@ -0,0 +1 @@
+from __future__ import absolute_import, division, print_function

pyro/contrib/gp/__init__.py

@@ -0,0 +1 @@
+from __future__ import absolute_import, division, print_function

pyro/contrib/gp/kernels/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import absolute_import, division, print_function
+
+from .rbf import RBF
+
+# flake8: noqa

pyro/contrib/gp/kernels/kernel.py

@@ -0,0 +1,49 @@
+from __future__ import absolute_import, division, print_function
+
+import torch.nn as nn
+
+
+class Kernel(nn.Module):
+    """
+    Base class for kernels used in Gaussian Process.
+
+    Every inherited class should implement the forward pass which
+        take inputs X, X2 and return their covariance matrix.
+    """
+
+    def __init__(self, active_dims=None, name=None):
+        super(Kernel, self).__init__()
+        self.active_dims = active_dims
+        self.name = name
+
+    def K(self, X, Z=None):
+        """
+        Calculate covariance matrix of inputs on active dimensionals.
+
+        :param torch.autograd.Variable X: A 2D tensor of size `N x input_dim`.
+        :param torch.autograd.Variable X2: A 2D tensor of size `N x input_dim`.
+        :return: Covariance matrix of X and X2 with size `N x N`.
+        :rtype: torch.autograd.Variable
+        """
+        if Z is None:
+            Z = X
+        X = self._slice_X(X)
+        Z = self._slice_X(Z)
+        K = self(X, Z)
+        return K
+
+    def _slice_X(self, X):
+        """
+        :param torch.autograd.Variable X: A 2D tensor.
+        :return: Slice X according to `self.active_dims`.
+        :rtype: torch.autograd.Variable
+        """
+        if X.dim() == 2:
+            active_dims = self.active_dims
+            if active_dims is None:
+                active_dims = slice(X.size(1))
+            return X[:, active_dims]
+        elif X.dim() == 1:
+            return X.unsqueeze(1)
+        else:
+            raise ValueError("Input X must be either 1 or 2 dimensional.")

pyro/contrib/gp/kernels/rbf.py

@@ -0,0 +1,33 @@
+from __future__ import absolute_import, division, print_function
+
+import torch
+from torch.nn import Parameter
+
+from .kernel import Kernel
+
+
+class RBF(Kernel):
+    """
+    Implementation of RBF kernel.
+    """
+
+    def __init__(self, variance=torch.ones(1), lengthscale=torch.ones(1), active_dims=None, name="RBF"):
+        super(RBF, self).__init__(active_dims=active_dims, name=name)
+        self.variance = Parameter(variance)
+        self.lengthscale = Parameter(lengthscale)
+
+    def forward(self, X, Z=None):
+        if Z is None:
+            Z = X
+        if X.dim() == 1:
+            X = X.unsqueeze(1)
+        if Z.dim() == 1:
+            Z = Z.unsqueeze(1)
+        if X.size(1) != Z.size(1):
+            raise ValueError("Inputs must have the same number of features.")
+
+        X2 = (X ** 2).sum(1, keepdim=True)
+        Z2 = (Z ** 2).sum(1, keepdim=True)
+        XZ = X.matmul(Z.t())
+        d2 = X2 - 2 * XZ + Z2.t()
+        return self.variance * torch.exp(-0.5 * d2 / self.lengthscale**2)

pyro/contrib/gp/models/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import absolute_import, division, print_function
+
+from .gpr import GPRegression
+
+# flake8: noqa

pyro/contrib/gp/models/gpr.py

@@ -0,0 +1,63 @@
+from __future__ import absolute_import, division, print_function
+
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+
+import pyro
+import pyro.distributions as dist
+
+
+class GPRegression(nn.Module):
+    """
+    Gaussian Process regression module.
+    """
+    def __init__(self, X, y, kernel, noise=torch.ones(1), priors={}):
+        super(GPRegression, self).__init__()
+        self.X = X
+        self.y = y
+        self.input_dim = X.size(0)
+        self.kernel = kernel
+        # TODO: define noise as a nn.Module, so we can train/set prior to it
+        self.noise = Variable(noise.type_as(X.data))
+        self.priors = priors
+
+    def model(self):
+        kernel_fn = pyro.random_module(self.kernel.name, self.kernel, self.priors)
+        kernel = kernel_fn()
+        K = kernel.K(self.X) + self.noise.repeat(self.input_dim).diag()
+        zero_loc = Variable(torch.zeros(self.input_dim).type_as(K.data))
+        pyro.sample("f", dist.MultivariateNormal(zero_loc, K), obs=self.y)
+
+    def guide(self):
+        guide_priors = {}
+        for p in self.priors:
+            p_MAP_name = pyro.param_with_module_name(self.kernel.name, p) + "_MAP"
+            # init params by their prior means
+            p_MAP = pyro.param(p_MAP_name, Variable(self.priors[p].analytic_mean().data.clone(),
+                                                    requires_grad=True))
+            guide_priors[p] = dist.Delta(p_MAP)
+        kernel_fn = pyro.random_module(self.kernel.name, self.kernel, guide_priors)
+        return kernel_fn()
+
+    def forward(self, Z):
+        """
+        Compute the parameters of `p(y|Z) ~ N(loc, covariance_matrix)`
+            w.r.t. the new input Z.
+
+        :param torch.autograd.Variable Z: A 2D tensor.
+        :return: loc and covariance matrix of p(y|Z).
+        :rtype: torch.autograd.Variable and torch.autograd.Variable
+        """
+        if Z.dim() == 2 and self.X.size(1) != Z.size(1):
+            assert ValueError("Train data and test data should have the same feature sizes.")
+        if Z.dim() == 1:
+            Z = Z.unsqueeze(1)
+        kernel = self.guide()
+        K = kernel.K(self.X) + self.noise.repeat(self.input_dim).diag()
+        K_xz = kernel(self.X, Z)
+        K_zx = K_xz.t()
+        K_zz = kernel.K(Z)
+        loc = K_zx.matmul(self.y.gesv(K)[0]).squeeze(1)
+        covariance_matrix = K_zz - K_zx.matmul(K_xz.gesv(K)[0])
+        return loc, covariance_matrix

tests/contrib/gp/__init__.py

@@ -0,0 +1 @@
+from __future__ import absolute_import, division, print_function

tests/contrib/gp/test_kernels.py

@@ -0,0 +1,18 @@
+from __future__ import absolute_import, division, print_function
+
+import torch
+from torch.autograd import Variable
+
+from pyro.contrib.gp.kernels import RBF
+from tests.common import assert_equal
+
+
+def test_K_rbf():
+    kernel = RBF(variance=torch.Tensor([2]), lengthscale=torch.Tensor([2]))
+    X = Variable(torch.Tensor([[1, 0, 1], [2, 1, 3]]))
+    Z = Variable(torch.Tensor([[4, 5, 6], [3, 1, 7]]))
+    K = kernel.K(X, Z)
+    assert K.dim() == 2
+    assert K.size(0) == 2
+    assert K.size(1) == 2
+    assert_equal(K.data.sum(), 0.30531)

tests/contrib/gp/test_models.py

@@ -0,0 +1,24 @@
+from __future__ import absolute_import, division, print_function
+
+import torch
+from torch.autograd import Variable
+
+from pyro.contrib.gp.kernels import RBF
+from pyro.contrib.gp.models import GPRegression
+from tests.common import assert_equal
+
+
+def test_forward_gpr():
+    kernel = RBF(torch.ones(1), torch.ones(1))
+    X = Variable(torch.Tensor([[1, 2, 3], [4, 5, 6]]))
+    y = Variable(torch.Tensor([0, 1]))
+    gpr = GPRegression(X, y, kernel, noise=torch.zeros(1))
+    Z = X
+    loc, cov = gpr(Z)
+    assert loc.dim() == 1
+    assert cov.dim() == 2
+    assert loc.size(0) == 2
+    assert cov.size(0) == 2
+    assert cov.size(1) == 2
+    assert_equal(loc.data.sum(), kernel.K(X).matmul(y).data.sum())
+    assert_equal(cov.data.abs().sum(), 0)