qpytorch 0.1__py3-none-any.whl

qpytorch/models/qeplvm/bayesian_qeplvm.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+
+from ..approximate_qep import ApproximateQEP
+
+
+class BayesianQEPLVM(ApproximateQEP):
+    """
+    The Q-Exponential Process Latent Variable Model (QEPLVM) class for unsupervised learning.
+    The class supports
+
+    1. Point estimates for latent X when prior_x = None
+    2. MAP Inference for X when prior_x is not None and inference == 'map'
+    3. Q-Exponential variational distribution q(X) when prior_x is not None and inference == 'variational'
+
+    .. seealso::
+        The `GPLVM tutorial
+        <examples/04_Variational_and_Approximate_GPs/Gaussian_Process_Latent_Variable_Models_with_Stochastic_Variational_Inference.ipynb>`_
+        for use instructions.
+
+    :param X: An instance of a sub-class of the LatentVariable class. One of,
+        :class:`~gpytorch.models.qeplvm.PointLatentVariable`, :class:`~gpytorch.models.qeplvm.MAPLatentVariable`, or
+        :class:`~gpytorch.models.qeplvm.VariationalLatentVariable`, to facilitate inference with 1, 2, or 3 respectively.
+    :type X: ~gpytorch.models.LatentVariable
+    :param ~gpytorch.variational._VariationalStrategy variational_strategy: The strategy that determines
+        how the model marginalizes over the variational distribution (over inducing points)
+        to produce the approximate posterior distribution (over data)
+    """
+
+    def __init__(self, X, variational_strategy):
+        super().__init__(variational_strategy)
+
+        # Assigning Latent Variable
+        self.X = X
+
+    def forward(self):
+        raise NotImplementedError
+
+    def sample_latent_variable(self):
+        sample = self.X()
+        return sample

qpytorch/models/qeplvm/latent_variable.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+
+import torch
+
+from ...module import Module
+
+
+class LatentVariable(Module):
+    """
+    This super class is used to describe the type of inference
+    used for the latent variable :math:`\\mathbf X` in QEPLVM models.
+
+    :param int n: Size of the latent space.
+    :param int latent_dim: Dimensionality of latent space.
+    """
+
+    def __init__(self, n, dim):
+        super().__init__()
+        self.n = n
+        self.latent_dim = dim
+
+    def forward(self, x):
+        raise NotImplementedError
+
+
+class PointLatentVariable(LatentVariable):
+    """
+    This class is used for QEPLVM models to recover a MLE estimate of
+    the latent variable :math:`\\mathbf X`.
+
+    :param int n: Size of the latent space.
+    :param int latent_dim: Dimensionality of latent space.
+    :param torch.Tensor X_init: initialization for the point estimate of :math:`\\mathbf X`
+    """
+
+    def __init__(self, n, latent_dim, X_init):
+        super().__init__(n, latent_dim)
+        self.register_parameter("X", X_init)
+
+    def forward(self):
+        return self.X
+
+
+class MAPLatentVariable(LatentVariable):
+    """
+    This class is used for QEPLVM models to recover a MAP estimate of
+    the latent variable :math:`\\mathbf X`, based on some supplied prior.
+
+    :param int n: Size of the latent space.
+    :param int latent_dim: Dimensionality of latent space.
+    :param torch.Tensor X_init: initialization for the point estimate of :math:`\\mathbf X`
+    :param ~gpytorch.priors.Prior prior_x: prior for :math:`\\mathbf X`
+    """
+
+    def __init__(self, n, latent_dim, X_init, prior_x):
+        super().__init__(n, latent_dim)
+        self.prior_x = prior_x
+        self.register_parameter("X", X_init)
+        self.register_prior("prior_x", prior_x, "X")
+
+    def forward(self):
+        return self.X
+
+
+class VariationalLatentVariable(LatentVariable):
+    """
+    This class is used for QEPLVM models to recover a variational approximation of
+    the latent variable :math:`\\mathbf X`. The variational approximation will be
+    an isotropic Q-Exponential distribution.
+
+    :param int n: Size of the latent space.
+    :param int data_dim: Dimensionality of the :math:`\\mathbf Y` values.
+    :param int latent_dim: Dimensionality of latent space.
+    :param torch.Tensor X_init: initialization for the point estimate of :math:`\\mathbf X`
+    :param ~gpytorch.priors.Prior prior_x: prior for :math:`\\mathbf X`
+    """
+
+    def __init__(self, n, data_dim, latent_dim, X_init, prior_x, **kwargs):
+        super().__init__(n, latent_dim)
+
+        self.data_dim = data_dim
+        self.prior_x = prior_x
+        # G: there might be some issues here if someone calls .cuda() on their BayesianQEPLVM
+        # after initializing on the CPU
+
+        # Local variational params per latent point with dimensionality latent_dim
+        self.q_mu = torch.nn.Parameter(X_init)
+        self.q_log_sigma = torch.nn.Parameter(torch.randn(n, latent_dim))
+        # This will add the KL divergence KL(q(X) || p(X)) to the loss
+        self.register_added_loss_term("x_kl")
+        
+        self.power = kwargs.pop('power', getattr(self.prior_x, 'power', torch.tensor(2.0)))
+
+    def forward(self):
+        from ...distributions import QExponential
+        from ...mlls import KLQExponentialAddedLossTerm
+
+        # Variational distribution over the latent variable q(x)
+        q_x = QExponential(self.q_mu, torch.nn.functional.softplus(self.q_log_sigma), power=self.power)
+        x_kl = KLQExponentialAddedLossTerm(q_x, self.prior_x, self.n, self.data_dim)
+        self.update_added_loss_term("x_kl", x_kl)  # Update the KL term
+        return q_x.rsample()

qpytorch/module.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+
+from gpytorch.module import Module as GModule
+
+
+class Module(GModule):
+    def named_hyperparameters(self):
+        from .variational._variational_distribution import _VariationalDistribution
+
+        for module_prefix, module in self.named_modules():
+            if not isinstance(module, _VariationalDistribution):
+                for elem in module.named_parameters(prefix=module_prefix, recurse=False):
+                    yield elem
+
+    def named_variational_parameters(self):
+        from .variational._variational_distribution import _VariationalDistribution
+
+        for module_prefix, module in self.named_modules():
+            if isinstance(module, _VariationalDistribution):
+                for elem in module.named_parameters(prefix=module_prefix, recurse=False):
+                    yield elem
+
+    def update_added_loss_term(self, name, added_loss_term):
+        from .mlls import AddedLossTerm
+
+        if not isinstance(added_loss_term, AddedLossTerm):
+            raise RuntimeError("added_loss_term must be a AddedLossTerm")
+        if name not in self._added_loss_terms.keys():
+            raise RuntimeError("added_loss_term {} not registered".format(name))
+        self._added_loss_terms[name] = added_loss_term

qpytorch/optim/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+from gpytorch.optim.ngd import NGD
+
+__all__ = ["NGD"]

qpytorch/priors/__init__.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+from gpytorch.priors.horseshoe_prior import HorseshoePrior
+from gpytorch.priors.lkj_prior import LKJCholeskyFactorPrior, LKJCovariancePrior, LKJPrior
+from gpytorch.priors.prior import Prior
+from gpytorch.priors.smoothed_box_prior import SmoothedBoxPrior
+from gpytorch.priors.torch_priors import (
+    GammaPrior,
+    HalfCauchyPrior,
+    HalfNormalPrior,
+    LogNormalPrior,
+    MultivariateNormalPrior,
+    NormalPrior,
+    UniformPrior,
+)
+from .qep_priors import (
+    MultivariateQExponentialPrior,
+    QExponentialPrior,
+)
+
+# from .wishart_prior import InverseWishartPrior, WishartPrior
+
+
+__all__ = [
+    "Prior",
+    "GammaPrior",
+    "HalfCauchyPrior",
+    "HalfNormalPrior",
+    "HorseshoePrior",
+    "LKJPrior",
+    "LKJCholeskyFactorPrior",
+    "LKJCovariancePrior",
+    "LogNormalPrior",
+    "MultivariateNormalPrior",
+    "MultivariateQExponentialPrior",
+    "NormalPrior",
+    "QExponentialPrior",
+    "SmoothedBoxPrior",
+    "UniformPrior",
+    # "InverseWishartPrior",
+    # "WishartPrior",
+]

qpytorch/priors/qep_priors.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+
+import torch
+from torch.nn import Module as TModule
+from linear_operator import to_linear_operator
+
+from ..distributions import QExponential, MultivariateQExponential
+from gpytorch.priors.prior import Prior
+from gpytorch.priors.utils import _bufferize_attributes, _del_attributes
+
+QEP_LAZY_PROPERTIES = ("covariance_matrix",)
+
+
+class QExponentialPrior(Prior, QExponential):
+    """
+    QExponential Prior
+
+    pdf(x) = q/2 * (2 * pi * sigma^2)^-0.5 * |(x - mu)/sigma|^(q/2-1) * exp(-0.5*|(x - mu)/sigma|^q)
+
+    where mu is the mean and sigma^2 is the variance.
+    """
+
+    def __init__(self, loc, scale, power=torch.tensor(1.0), validate_args=False, transform=None):
+        TModule.__init__(self)
+        QExponential.__init__(self, loc=loc, scale=scale, power=power, validate_args=validate_args)
+        _bufferize_attributes(self, ("loc", "scale"))
+        self._transform = transform
+
+    def expand(self, batch_shape):
+        batch_shape = torch.Size(batch_shape)
+        return QExponentialPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape), self.power)
+
+
+class MultivariateQExponentialPrior(Prior, MultivariateQExponential):
+    """Multivariate Q-Exponential prior
+
+    pdf(x) = q/2 * det(2 * pi * Sigma)^-0.5 * r^((q/2-1)*d/2) * exp(-0.5 * r^(q/2)), r = (x - mu)' Sigma^-1 (x - mu)
+
+    where mu is the mean and Sigma > 0 is the covariance matrix.
+    """
+
+    def __init__(
+        self, mean, covariance_matrix, power=torch.tensor(1.0), validate_args=False, transform=None
+    ):
+        TModule.__init__(self)
+        MultivariateQExponential.__init__(
+            self,
+            mean=mean,
+            covariance_matrix=covariance_matrix,
+            power=power,
+            validate_args=validate_args,
+        )
+        _bufferize_attributes(self, ("loc",))
+        self._transform = transform
+
+    def cuda(self, device=None):
+        """Applies module-level cuda() call and resets all lazy properties"""
+        module = self._apply(lambda t: t.cuda(device))
+        _del_attributes(module, QEP_LAZY_PROPERTIES)
+        return module
+
+    def cpu(self):
+        """Applies module-level cpu() call and resets all lazy properties"""
+        module = self._apply(lambda t: t.cpu())
+        _del_attributes(module, QEP_LAZY_PROPERTIES)
+        return module
+
+    @property
+    def lazy_covariance_matrix(self):
+        if self.islazy:
+            return self._covar
+        else:
+            return to_linear_operator(super().covariance_matrix)
+
+    def expand(self, batch_shape):
+        batch_shape = torch.Size(batch_shape)
+        cov_shape = batch_shape + self.event_shape
+        new_loc = self.loc.expand(batch_shape)
+        new_covar = self._covar.expand(cov_shape)
+
+        return MultivariateQExponentialPrior(mean=new_loc, covariance_matrix=new_covar, power=self.power)

qpytorch/test/__init__.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+
+from gpytorch.test.base_test_case import BaseTestCase
+from gpytorch.test.base_keops_test_case import BaseKeOpsTestCase
+from gpytorch.test.base_kernel_test_case import BaseKernelTestCase
+from .base_likelihood_test_case import BaseLikelihoodTestCase
+from gpytorch.test.base_mean_test_case import BaseMeanTestCase
+from .model_test_case import BaseModelTestCase, VariationalModelTestCase
+from gpytorch.test import utils
+from .variational_test_case import VariationalTestCase
+
+__all__ = [
+    "BaseKeOpsTestCase",
+    "BaseKernelTestCase",
+    "BaseLikelihoodTestCase",
+    "BaseMeanTestCase",
+    "BaseModelTestCase",
+    "BaseTestCase"
+    "utils",
+    "VariationalModelTestCase",
+    "VariationalTestCase",
+]

qpytorch/test/base_likelihood_test_case.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+
+from abc import abstractmethod
+
+import torch
+from torch.distributions import Distribution
+
+import qpytorch
+from ..distributions import MultivariateNormal, MultivariateQExponential
+from ..likelihoods import Likelihood
+
+from gpytorch.test.base_test_case import BaseTestCase
+
+
+class BaseLikelihoodTestCase(BaseTestCase):
+    @abstractmethod
+    def create_likelihood(self, **kwargs):
+        raise NotImplementedError()
+
+    def _create_conditional_input(self, batch_shape=torch.Size()):
+        return torch.randn(*batch_shape, 5)
+
+    def _create_marginal_input(self, batch_shape=torch.Size()):
+        mat = torch.randn(*batch_shape, 5, 5)
+        eye = torch.diag_embed(torch.ones(*batch_shape, 5))
+        if 'Gaussian' in self.__class__.__name__:
+            return MultivariateNormal(torch.randn(*batch_shape, 5), mat @ mat.transpose(-1, -2) + eye)
+        elif 'QExponential' in self.__class__.__name__:
+            return MultivariateQExponential(torch.randn(*batch_shape, 5), mat @ mat.transpose(-1, -2) + eye, torch.tensor(getattr(self, '_power', 2.0)))
+
+    def _create_targets(self, batch_shape=torch.Size()):
+        return torch.randn(*batch_shape, 5)
+
+    def _test_conditional(self, batch_shape):
+        likelihood = self.create_likelihood()
+        likelihood.max_plate_nesting += len(batch_shape)
+        input = self._create_conditional_input(batch_shape)
+        output = likelihood(input)
+
+        self.assertTrue(isinstance(output, Distribution))
+        self.assertEqual(output.sample().shape, input.shape)
+
+    def _test_log_marginal(self, batch_shape):
+        likelihood = self.create_likelihood()
+        likelihood.max_plate_nesting += len(batch_shape)
+        input = self._create_marginal_input(batch_shape)
+        target = self._create_targets(batch_shape)
+        with qpytorch.settings.num_likelihood_samples(512):
+            output = likelihood.log_marginal(target, input)
+
+        self.assertTrue(torch.is_tensor(output))
+        self.assertEqual(output.shape, batch_shape + torch.Size([5]))
+        with qpytorch.settings.num_likelihood_samples(512):
+            default_log_prob = Likelihood.log_marginal(likelihood, target, input)
+        self.assertAllClose(output, default_log_prob, rtol=0.25)
+
+    def _test_log_prob(self, batch_shape):
+        likelihood = self.create_likelihood()
+        likelihood.max_plate_nesting += len(batch_shape)
+        input = self._create_marginal_input(batch_shape)
+        target = self._create_targets(batch_shape)
+        with qpytorch.settings.num_likelihood_samples(512):
+            output = likelihood.expected_log_prob(target, input)
+
+        self.assertTrue(torch.is_tensor(output))
+        self.assertEqual(output.shape, batch_shape + torch.Size([5]))
+        with qpytorch.settings.num_likelihood_samples(512):
+            default_log_prob = Likelihood.expected_log_prob(likelihood, target, input)
+        self.assertAllClose(output, default_log_prob, rtol=0.25)
+
+    def _test_marginal(self, batch_shape):
+        likelihood = self.create_likelihood()
+        likelihood.max_plate_nesting += len(batch_shape)
+        input = self._create_marginal_input(batch_shape)
+        output = likelihood(input)
+
+        self.assertTrue(isinstance(output, Distribution))
+        self.assertEqual(output.sample().shape[-len(input.sample().shape) :], input.sample().shape)
+
+        # Compare against default implementation
+        with qpytorch.settings.num_likelihood_samples(30000):
+            default = Likelihood.marginal(likelihood, input)
+        # print(output.mean, default.mean)
+        default_mean = default.mean
+        actual_mean = output.mean
+        if default_mean.dim() > actual_mean.dim():
+            default_mean = default_mean.mean(0)
+        self.assertAllClose(default_mean, actual_mean, rtol=0.25, atol=0.25)
+
+    def test_nonbatch(self):
+        self._test_conditional(batch_shape=torch.Size([]))
+        self._test_log_marginal(batch_shape=torch.Size([]))
+        self._test_log_prob(batch_shape=torch.Size([]))
+        self._test_marginal(batch_shape=torch.Size([]))
+
+    def test_batch(self):
+        self._test_conditional(batch_shape=torch.Size([3]))
+        self._test_log_marginal(batch_shape=torch.Size([3]))
+        self._test_log_prob(batch_shape=torch.Size([3]))
+        self._test_marginal(batch_shape=torch.Size([3]))
+
+    def test_multi_batch(self):
+        self._test_conditional(batch_shape=torch.Size([2, 3]))
+        self._test_log_marginal(batch_shape=torch.Size([2, 3]))
+        self._test_log_prob(batch_shape=torch.Size([2, 3]))
+        self._test_marginal(batch_shape=torch.Size([2, 3]))

qpytorch/test/model_test_case.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+
+from abc import abstractmethod
+
+import torch
+
+import qpytorch
+
+
+class BaseModelTestCase(object):
+    @abstractmethod
+    def create_model(self, train_x, train_y, likelihood):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_test_data(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_likelihood_and_labels(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_batch_test_data(self, batch_shape=torch.Size([3])):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_batch_likelihood_and_labels(self, batch_shape=torch.Size([3])):
+        raise NotImplementedError()
+
+    def test_forward_train(self):
+        data = self.create_test_data()
+        likelihood, labels = self.create_likelihood_and_labels()
+        model = self.create_model(data, labels, likelihood)
+        model.train()
+        output = model(data)
+        self.assertTrue(output.lazy_covariance_matrix.dim() == 2)
+        self.assertTrue(output.lazy_covariance_matrix.size(-1) == data.size(-2))
+        self.assertTrue(output.lazy_covariance_matrix.size(-2) == data.size(-2))
+
+    def test_batch_forward_train(self):
+        batch_data = self.create_batch_test_data()
+        likelihood, labels = self.create_batch_likelihood_and_labels()
+        model = self.create_model(batch_data, labels, likelihood)
+        model.train()
+        output = model(batch_data)
+        self.assertTrue(output.lazy_covariance_matrix.dim() == 3)
+        self.assertTrue(output.lazy_covariance_matrix.size(-1) == batch_data.size(-2))
+        self.assertTrue(output.lazy_covariance_matrix.size(-2) == batch_data.size(-2))
+
+    def test_multi_batch_forward_train(self):
+        batch_data = self.create_batch_test_data(batch_shape=torch.Size([2, 3]))
+        likelihood, labels = self.create_batch_likelihood_and_labels(batch_shape=torch.Size([2, 3]))
+        model = self.create_model(batch_data, labels, likelihood)
+        model.train()
+        output = model(batch_data)
+        self.assertTrue(output.lazy_covariance_matrix.dim() == 4)
+        self.assertTrue(output.lazy_covariance_matrix.size(-1) == batch_data.size(-2))
+        self.assertTrue(output.lazy_covariance_matrix.size(-2) == batch_data.size(-2))
+
+    def test_forward_eval(self):
+        data = self.create_test_data()
+        likelihood, labels = self.create_likelihood_and_labels()
+        model = self.create_model(data, labels, likelihood)
+        model.eval()
+        output = model(data)
+        self.assertTrue(output.lazy_covariance_matrix.dim() == 2)
+        self.assertTrue(output.lazy_covariance_matrix.size(-1) == data.size(-2))
+        self.assertTrue(output.lazy_covariance_matrix.size(-2) == data.size(-2))
+
+    def test_batch_forward_eval(self):
+        batch_data = self.create_batch_test_data()
+        likelihood, labels = self.create_batch_likelihood_and_labels()
+        model = self.create_model(batch_data, labels, likelihood)
+        model.eval()
+        output = model(batch_data)
+        self.assertTrue(output.lazy_covariance_matrix.dim() == 3)
+        self.assertTrue(output.lazy_covariance_matrix.size(-1) == batch_data.size(-2))
+        self.assertTrue(output.lazy_covariance_matrix.size(-2) == batch_data.size(-2))
+
+    def test_multi_batch_forward_eval(self):
+        batch_data = self.create_batch_test_data(batch_shape=torch.Size([2, 3]))
+        likelihood, labels = self.create_batch_likelihood_and_labels(batch_shape=torch.Size([2, 3]))
+        model = self.create_model(batch_data, labels, likelihood)
+        model.eval()
+        output = model(batch_data)
+        self.assertTrue(output.lazy_covariance_matrix.dim() == 4)
+        self.assertTrue(output.lazy_covariance_matrix.size(-1) == batch_data.size(-2))
+        self.assertTrue(output.lazy_covariance_matrix.size(-2) == batch_data.size(-2))
+
+
+class VariationalModelTestCase(BaseModelTestCase):
+    def test_backward_train(self):
+        data = self.create_test_data()
+        likelihood, labels = self.create_likelihood_and_labels()
+        model = self.create_model(data, labels, likelihood)
+        mll = qpytorch.mlls.VariationalELBO(likelihood, model, num_data=labels.size(-1))
+        model.train()
+        likelihood.train()
+
+        # We'll just do one step of gradient descent to mix up the params a bit
+        optimizer = torch.optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
+
+        output = model(data)
+        loss = -mll(output, labels)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        output = model(data)
+        loss = -mll(output, labels)
+        loss.backward()
+
+        for _, param in model.named_parameters():
+            self.assertTrue(param.grad is not None)
+            self.assertGreater(param.grad.norm().item(), 0)
+        for _, param in likelihood.named_parameters():
+            self.assertTrue(param.grad is not None)
+            self.assertGreater(param.grad.norm().item(), 0)
+        optimizer.step()
+
+    def test_batch_backward_train(self, batch_shape=torch.Size([3])):
+        data = self.create_batch_test_data(batch_shape)
+        likelihood, labels = self.create_batch_likelihood_and_labels(batch_shape)
+        model = self.create_model(data, labels, likelihood)
+        mll = qpytorch.mlls.VariationalELBO(likelihood, model, num_data=labels.size(-1))
+        model.train()
+        likelihood.train()
+
+        # We'll just do one step of gradient descent to mix up the params a bit
+        optimizer = torch.optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
+
+        output = model(data)
+        loss = -mll(output, labels).sum()
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        output = model(data)
+        loss = -mll(output, labels).sum()
+        loss.backward()
+
+        for _, param in model.named_parameters():
+            self.assertTrue(param.grad is not None)
+            self.assertGreater(param.grad.norm().item(), 0)
+        for _, param in likelihood.named_parameters():
+            self.assertTrue(param.grad is not None)
+            self.assertGreater(param.grad.norm().item(), 0)
+        optimizer.step()
+
+    def test_multi_batch_backward_train(self, batch_shape=torch.Size([2, 3])):
+        return self.test_batch_backward_train(batch_shape=batch_shape)