Gaussian Process tutorial

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, input_dim, active_dims=None, name=None):
+        super(Kernel, self).__init__()
+        if active_dims is None:
+            active_dims = slice(input_dim)
+        elif input_dim != len(active_dims):
+            raise ValueError("Input size and the length of active dimensionals should be equal.")
+        self.input_dim = input_dim
+        self.active_dims = active_dims
+        self.name = name
+
+    def forward(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 Z: A 2D tensor of size `N x input_dim`.
+        :return: Covariance matrix of X and Z with size `N x N`.
+        :rtype: torch.autograd.Variable
+        """
+        raise NotImplementedError
+
+    def _slice_X(self, X):
+        """
+        Slice X according to `self.active_dims`. If X is 1 dimensional then returns
+            a 2D tensor of size `N x 1`.
+
+        :param torch.autograd.Variable X: A 1D or 2D tensor.
+        :return: A 2D slice of X.
+        :rtype: torch.autograd.Variable
+        """
+        if X.dim() == 2:
+            return X[:, self.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,46 @@
+from __future__ import absolute_import, division, print_function
+
+import torch
+from torch.nn import Parameter
+
+from .kernel import Kernel
+
+
+class RBF(Kernel):
+    """
+    Implementation of Radial Basis Function kernel.
+
+    By default, parameters will be `torch.nn.Parameter`s containing `torch.FloatTensor`s.
+        To cast them to the correct data type or GPU device, we can call methods such as
+        `.double()`, `.cuda(device=None)`,... See
+        `torch.nn.Module <http://pytorch.org/docs/master/nn.html#torch.nn.Module>`_ for more information.
+
+    :param int input_dim: Dimension of inputs for this kernel.
+    :param torch.Tensor variance: Variance parameter of this kernel.
+    :param torch.Tensor lengthscale: Length scale parameter of this kernel.
+    """
+
+    def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None, name="RBF"):
+        super(RBF, self).__init__(input_dim, active_dims, name)
+        if variance is None:
+            variance = torch.ones(1)
+        self.variance = Parameter(variance)
+        if lengthscale is None:
+            lengthscale = torch.ones(input_dim)
+        self.lengthscale = Parameter(lengthscale)
+
+    def forward(self, X, Z=None):
+        if Z is None:
+            Z = X
+        X = self._slice_X(X)
+        Z = self._slice_X(Z)
+        if X.size(1) != Z.size(1):
+            raise ValueError("Inputs must have the same number of features.")
+
+        scaled_X = X / self.lengthscale
+        scaled_Z = Z / self.lengthscale
+        X2 = (scaled_X ** 2).sum(1, keepdim=True)
+        Z2 = (scaled_Z ** 2).sum(1, keepdim=True)
+        XZ = scaled_X.matmul(scaled_Z.t())
+        d2 = X2 - 2 * XZ + Z2.t()
+        return self.variance * torch.exp(-0.5 * d2)

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,76 @@
+from __future__ import absolute_import, division, print_function
+
+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.
+
+    :param torch.autograd.Variable X: A tensor of inputs.
+    :param torch.autograd.Variable y: A tensor of outputs for training.
+    :param pyro.gp.kernels.Kernel kernel: A Pyro kernel object.
+    :param torch.Tensor noise: An optional noise tensor.
+    :param dict priors: A mapping from kernel parameter's names to priors.
+    """
+    def __init__(self, X, y, kernel, noise=None, priors=None):
+        super(GPRegression, self).__init__()
+        self.X = X
+        self.y = y
+        self.input_dim = X.size(0)
+        self.kernel = kernel
+        # TODO: define noise as a Likelihood (another nn.Module beside kernel),
+        # then we can train/set prior to it
+        if noise is None:
+            self.noise = Variable(X.data.new([1]))
+        else:
+            self.noise = Variable(noise)
+        self.priors = priors
+        if priors is None:
+            self.priors = {}
+
+    def model(self):
+        kernel_fn = pyro.random_module(self.kernel.name, self.kernel, self.priors)
+        kernel = kernel_fn()
+        K = kernel(self.X) + self.noise.repeat(self.input_dim).diag()
+        zero_loc = Variable(K.data.new([0]).expand(self.input_dim))
+        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(self.X) + self.noise.repeat(self.input_dim).diag()
+        K_xz = kernel(self.X, Z)
+        K_zx = K_xz.t()
+        K_zz = kernel(Z)
+        # TODO: use torch.trtrs when it supports cuda tensors
+        K_inv = K.inverse()
+        loc = K_zx.matmul(K_inv.matmul(self.y))
+        covariance_matrix = K_zz - K_zx.matmul(K_inv.matmul(K_xz))
+        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_forward_rbf():
+    kernel = RBF(input_dim=3, variance=torch.Tensor([2]), lengthscale=torch.Tensor([2, 2, 2]))
+    X = Variable(torch.Tensor([[1, 0, 1], [2, 1, 3]]))
+    Z = Variable(torch.Tensor([[4, 5, 6], [3, 1, 7]]))
+    K = kernel(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(input_dim=3, variance=torch.ones(1), lengthscale=torch.ones(3))
+    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(X).matmul(y).data.sum())
+    assert_equal(cov.data.abs().sum(), 0)