qpytorch 0.1__py3-none-any.whl

qpytorch/models/deep_qeps/deep_qep.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+
+import warnings
+
+import torch
+from linear_operator.operators import BlockDiagLinearOperator
+
+from ... import settings
+from ...distributions import QExponential, MultitaskMultivariateQExponential
+from ...likelihoods import Likelihood
+from ..approximate_qep import ApproximateQEP
+from ..qep import QEP
+
+
+class _DeepQEPVariationalStrategy(object):
+    def __init__(self, model):
+        self.model = model
+
+    @property
+    def sub_variational_strategies(self):
+        if not hasattr(self, "_sub_variational_strategies_memo"):
+            self._sub_variational_strategies_memo = [
+                module.variational_strategy for module in self.model.modules() if isinstance(module, ApproximateQEP)
+            ]
+        return self._sub_variational_strategies_memo
+
+    def kl_divergence(self):
+        return sum(strategy.kl_divergence().sum() for strategy in self.sub_variational_strategies)
+
+
+class DeepQEPLayer(ApproximateQEP):
+    """
+    Represents a layer in a deep QEP where inference is performed via the doubly stochastic method of
+    Salimbeni et al., 2017. Upon calling, instead of returning a variational distribution q(f), returns samples
+    from the variational distribution.
+
+    See the documentation for __call__ below for more details below. Note that the behavior of __call__
+    will change to be much more elegant with multiple batch dimensions; however, the interface doesn't really
+    change.
+
+    :param ~gpytorch.variational.VariationalStrategy variational_strategy: Strategy for
+        changing q(u) -> q(f) (see other VI docs)
+    :param int input_dims`: Dimensionality of input data expected by each QEP
+    :param int output_dims: (default None) Number of QEPs in this layer, equivalent to
+        output dimensionality. If set to `None`, then the output dimension will be squashed.
+
+    Forward data through this hidden QEP layer. The output is a MultitaskMultivariateQExponential distribution
+    (or MultivariateQExponential distribution is output_dims=None).
+
+    If the input is >=2 dimensional Tensor (e.g. `n x d`), we pass the input through each hidden QEP,
+    resulting in a `n x h` multitask Q-Exponential distribution (where all of the `h` tasks represent an
+    output dimension and are independent from one another).  We then draw `s` samples from these Q-Exponentials,
+    resulting in a `s x n x h` MultitaskMultivariateQExponential distribution.
+
+    If the input is a >=3 dimensional Tensor, and the `are_samples=True` kwarg is set, then we assume that
+    the outermost batch dimension is a samples dimension. The output will have the same number of samples.
+    For example, a `s x b x n x d` input will result in a `s x b x n x h` MultitaskMultivariateQExponential distribution.
+
+    The goal of these last two points is that if you have a tensor `x` that is `n x d`, then
+
+        >>> hidden_qep2(hidden_qep(x))
+
+    will just work, and return a tensor of size `s x n x h2`, where `h2` is the output dimensionality of
+    hidden_qep2. In this way, hidden QEP layers are easily composable.
+    """
+
+    def __init__(self, variational_strategy, input_dims, output_dims):
+        super(DeepQEPLayer, self).__init__(variational_strategy)
+        self.input_dims = input_dims
+        self.output_dims = output_dims
+
+    def forward(self, x):
+        raise NotImplementedError
+
+    def __call__(self, inputs, are_samples=False, **kwargs):
+        deterministic_inputs = not are_samples
+        if isinstance(inputs, MultitaskMultivariateQExponential):
+            inputs = QExponential(loc=inputs.mean, scale=inputs.variance.sqrt(), power=inputs.power).rsample(rescale=kwargs.pop('rescale', False))
+            deterministic_inputs = False
+
+        if settings.debug.on():
+            if not torch.is_tensor(inputs):
+                raise ValueError(
+                    "`inputs` should either be a MultitaskMultivariateQExponential or a Tensor, got "
+                    f"{inputs.__class__.__Name__}"
+                )
+
+            if inputs.size(-1) != self.input_dims:
+                raise RuntimeError(
+                    f"Input shape did not match self.input_dims. Got total feature dims [{inputs.size(-1)}],"
+                    f" expected [{self.input_dims}]"
+                )
+
+        # Repeat the input for all possible outputs
+        if self.output_dims is not None:
+            inputs = inputs.unsqueeze(-3)
+            inputs = inputs.expand(*inputs.shape[:-3], self.output_dims, *inputs.shape[-2:])
+
+        # Now run samples through the QEP
+        output = ApproximateQEP.__call__(self, inputs, **kwargs)
+        if self.output_dims is not None:
+            mean = output.loc.transpose(-1, -2)
+            covar = BlockDiagLinearOperator(output.lazy_covariance_matrix, block_dim=-3)
+            output = MultitaskMultivariateQExponential(mean, covar, power=output.power, interleaved=False)
+
+        # Maybe expand inputs?
+        if deterministic_inputs:
+            output = output.expand(torch.Size([settings.num_likelihood_samples.value()]) + output.batch_shape)
+
+        return output
+
+
+class DeepQEP(QEP):
+    """
+    A container module to build a DeepQEP.
+    This module should contain :obj:`~gpytorch.models.deep.DeepQEPLayer`
+    modules, and can also contain other modules as well.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.variational_strategy = _DeepQEPVariationalStrategy(self)
+
+    def forward(self, x):
+        raise NotImplementedError
+
+
+class DeepLikelihood(Likelihood):
+    """
+    A wrapper to make a GPyTorch likelihood compatible with Deep QEPs
+
+    Example:
+        >>> deep_qexponential_likelihood = gpytorch.likelihoods.DeepLikelihood(gpytorch.likelihood.QExponentialLikelihood)
+    """
+
+    def __init__(self, base_likelihood):
+        super().__init__()
+        warnings.warn(
+            "DeepLikelihood is now deprecated. Use a standard likelihood in conjunction with a "
+            "gpytorch.mlls.DeepApproximateMLL. See the DeepQEP example in our documentation.",
+            DeprecationWarning,
+        )
+        self.base_likelihood = base_likelihood
+
+    def expected_log_prob(self, observations, function_dist, *params, **kwargs):
+        return self.base_likelihood.expected_log_prob(observations, function_dist, *params, **kwargs).mean(dim=0)
+
+    def log_marginal(self, observations, function_dist, *params, **kwargs):
+        return self.base_likelihood.log_marginal(observations, function_dist, *params, **kwargs).mean(dim=0)
+
+    def forward(self, *args, **kwargs):
+        pass
+
+    def __call__(self, *args, **kwargs):
+        return self.base_likelihood.__call__(*args, **kwargs)

qpytorch/models/deep_qeps/dspp.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+import torch
+from linear_operator.operators import BlockDiagLinearOperator
+
+from ... import settings
+from ...distributions import MultitaskMultivariateQExponential
+from ..approximate_qep import ApproximateQEP
+from .deep_qep import DeepQEP, DeepQEPLayer
+
+
+class DSPPLayer(DeepQEPLayer):
+    """
+    Represents a layer in a DSPP where inference is performed using the techniques defined in Jankowiak et al., 2020.
+
+    From an end user perspective, the functionality and usage of a DSPPLayer is essentially identical to that of a
+    DeepQEPLayer. It is therefore recommended that you review the documentation for DeepQEPLayer.
+
+    :param ~gpytorch.variational.VariationalStrategy variational_strategy: Strategy for
+        changing q(u) -> q(f) (see other VI docs)
+    :param int input_dims: Dimensionality of input data expected by each QEP
+    :param int output_dims: (default None) Number of QEPs in this layer, equivalent to
+        output dimensionality. If set to `None`, then the output dimension will be squashed.
+    :param int num_quad_sites: Number of quadrature sites to use. Also the number of Q-Exponentials in the mixture output
+        by this layer.
+
+    Again, refer to the documentation for DeepQEPLayer or our example notebooks for full details on what calling a
+    DSPPLayer module does. The high level overview is that if a tensor `x` is `n x d` then
+
+        >>> hidden_qep2(hidden_qep1(x))
+
+    will return a `num_quad_sites` by `output_dims` set of Q-Exponentials, where for each output dim the first batch dim
+    represents a weighted mixture of `num_quad_sites` Q-Exponentials with weights given by DSPP.quad_weights (see DSPP below)
+    """
+
+    def __init__(self, variational_strategy, input_dims, output_dims, num_quad_sites=3, quad_sites=None):
+        super().__init__(variational_strategy, input_dims, output_dims)
+
+        self.num_quad_sites = num_quad_sites
+
+        # Pass in previous_layer.quad_sites if you want to share quad_sites across layers.
+        if quad_sites is not None:
+            self.quad_sites = quad_sites
+        else:
+            self.quad_sites = torch.nn.Parameter(torch.randn(num_quad_sites, input_dims))
+
+    def __call__(self, inputs, **kwargs):
+        if isinstance(inputs, MultitaskMultivariateQExponential):
+            # This is for subsequent layers. We apply quadrature here
+            # Mean, stdv are q x ... x n x t
+            mus, sigmas = inputs.mean, inputs.variance.sqrt()
+            qg = self.quad_sites.view([self.num_quad_sites] + [1] * (mus.dim() - 2) + [self.input_dims])
+            sigmas = sigmas * qg
+            inputs = mus + sigmas  # q^t x n x t
+            deterministic_inputs = False
+        else:
+            deterministic_inputs = True
+
+        if settings.debug.on():
+            if not torch.is_tensor(inputs):
+                raise ValueError(
+                    "`inputs` should either be a MultitaskMultivariateQExponential or a Tensor, got "
+                    f"{inputs.__class__.__Name__}"
+                )
+
+            if inputs.size(-1) != self.input_dims:
+                raise RuntimeError(
+                    f"Input shape did not match self.input_dims. Got total feature dims [{inputs.size(-1)}],"
+                    f" expected [{self.input_dims}]"
+                )
+
+        # Repeat the input for all possible outputs
+        if self.output_dims is not None:
+            inputs = inputs.unsqueeze(-3)
+            inputs = inputs.expand(*inputs.shape[:-3], self.output_dims, *inputs.shape[-2:])
+
+        # Now run samples through the QEP
+        output = ApproximateQEP.__call__(self, inputs, **kwargs)
+
+        # If this is the first layer (deterministic inputs), expand the output
+        # This allows quadrature to be applied to future layers
+        if deterministic_inputs:
+            output = output.expand(torch.Size([self.num_quad_sites]) + output.batch_shape)
+
+        if self.num_quad_sites > 0:
+            if self.output_dims is not None and not isinstance(output, MultitaskMultivariateQExponential):
+                mean = output.loc.transpose(-1, -2)
+                covar = BlockDiagLinearOperator(output.lazy_covariance_matrix, block_dim=-3)
+                output = MultitaskMultivariateQExponential(mean, covar, power=output.power, interleaved=False)
+        else:
+            output = output.loc.transpose(-1, -2)  # this layer provides noiseless kernel interpolation
+
+        return output
+
+
+class DSPP(DeepQEP):
+    """
+    A container module to build a DSPP
+    This module should contain :obj:`~gpytorch.models.deep_qeps.DSPPLayer`
+    modules, and can also contain other modules as well.
+
+    This Module contains an additional set of parameters, `raw_quad_weights`, that represent the mixture weights for
+    the output distribution.
+    """
+
+    def __init__(self, num_quad_sites):
+        super().__init__()
+        self.num_quad_sites = num_quad_sites
+        self.register_parameter("raw_quad_weights", torch.nn.Parameter(torch.randn(self.num_quad_sites)))
+
+    @property
+    def quad_weights(self):
+        qwd = self.raw_quad_weights
+        return qwd - qwd.logsumexp(dim=-1)