Repository: cornellius-gp/gpytorch

Issue:
[Bug] Fully Bayesian Inference with Pyro Does Not Work
# 🐛 Bug

Right now, the way our Bayesian inference with pyro works is that we loop over all `Parameter`s with priors, call `pyro.sample` and load the sample into the parameter value. Here's an extremely simplified model that is equivalent to what we do that does NOT work:
```python
def pyro_model2(x, y):
    raw_lengthscale = pyro.sample('raw_lengthscale', Normal(0, 1))
    raw_outputscale = pyro.sample('raw_outputscale', Normal(0, 1))

    # This is not an okay way to load the results of pyro.sample statements into a Module
    covar_module.initialize(**{'raw_outputscale': raw_outputscale, 'base_kernel.raw_lengthscale': raw_lengthscale})
    
    covar_x = covar_module(x).evaluate()
    pyro.sample("obs", pyro.distributions.MultivariateNormal(torch.zeros(y.size(-1)), covar_x + torch.eye(y.size(-1))), obs=y)
```
See the attached python notebook that demonstrates that this leads to completely wrong gradients for the potential energy, so HMC is entirely broken.

The best way I've found to fix this is to use the (possibly deprecated?) `pyro.random_module` primitive. Here's a pyro model using a full GPyTorch GP that gets correct potential derivatives:
```python
def pyro_model5(x, y):
    priors= {
        'covar_module.base_kernel.raw_lengthscale': Normal(0, 1).expand([1, 1]),
        'covar_module.raw_outputscale': Normal(0, 1),
        'likelihood.noise_covar.raw_noise': Normal(0, 1).expand([1]),
        'mean_module.constant': Normal(0, 1),
   }
    fn = pyro.random_module("model", model, prior=priors)
    sampled_model = fn()
    
    output = sampled_model.likelihood(sampled_model(x))
    pyro.sample("obs", output, obs=y)
```
Now, we could definitely just replace our existing GPyTorch interface to basically do the above instead for any parameter that has a prior registered. The only problem is that I don't know if this will let us place priors over functions of `Parameters` (i.e., if we want to place a prior over the `lengthscale` rather than the `raw_lengthscale`).

Thoughts? I feel like the above is *almost* there, but it'd be great to still support placing priors over derived values of the parameters.

cc/ @rmgarnett

also cc/ @fritzo @eb8680 @martinjankowiak in case any of the pyro devs hopefully have a thought about how to do this properly that I'm just missing?

[Sampling Bug.ipynb.txt](https://github.com/cornellius-gp/gpytorch/files/6328496/Sampling.Bug.ipynb.txt)


Discussion:
Cc @jpchen @Balandat @dme65 
I was looking into this. I think the issue here is that HMC/NUTS in Pyro expects a functional, i.e. one needs to be able to take gradients with respect to the arguments of the potential function. In this case, or more generally, when `nn.Module`s are involved and we replace the `nn.Parameter` values in the modules by replacing the `data` attribute, the graph gets disconnected, because `nn.Parameter`s are supposed to be leaf variables in the autograd graph. This should throw a `RuntimeError` which I suspect is being swallowed, likely somewhere here: https://github.com/cornellius-gp/gpytorch/blob/481f2856092681908d046acc1263aa0bc5f86877/gpytorch/module.py#L106.

This is the reason why snippet 1 does not work. As to why this works for the deprecated `pyro.random_module` is that it replaces these `nn.Parameter`s with actual tensors. While this might work for this use case, it has other unexpected side-effects, and is is one of the reasons why we do not encourage its broad use. If this works for your use for the time being, you can also replicate the same effect by using the following method in gpytorch's `Module` class to reset the parameters of the kernel (instead of calling `initialize`). This is exactly what `random_module` is doing under the hood (hacky!):

```python
  def reset_params_pyro(self, **params):
      for name, value in params.items():
          # find the innermost module
          names = name.rsplit(".")
          if len(names) > 1:
              mod_name, param_name = names
              mod = operator.attrgetter(mod_name)(self)
          else:
              mod, param_name = self, name
          # delete paramater (otherwise we get a TypeError when replacing Parameter -> Tensor)
          old_param = getattr(mod, param_name)
          delattr(mod, param_name)
          # XXX: replaces nn.Param by torch.tensor
          setattr(mod, param_name, value.expand(old_param.shape))
```

This mismatch is not an uncommon case and I wonder if something like https://github.com/facebookresearch/higher could be used to functionalize the module objects in a way that works with HMC. But the easier (and certain) route would be to build better support for `nn.Module` classes within Pyro's HMC, e.g. if we know that the PE function is a `nn.Module` we can dispatch it to a separate `potential_grad` function that is specialized for `nn.Module`.
@neerajprad Thanks! This is actually super helpful. At a minimum, the `reset_params_pyro` solution is much better for us than directly using `pyro.random_module` because it's obvious how it allows us to continue placing priors over transformed versions of the Parameters, and is only equally hacky.

Based on what you've said, I'm assuming the following workflow would work:

1. deepcopy the module.
2. Call `pyro.sample` for each prior and apply the inverse transform to get the appropriate value for the underlying raw parameter, as we do now.
3. Use `reset_params_pyro` above to replace the parameters with the samples on the copied module.

If so, I will probably implement things this way in the short term. It keeps most or all of the interface we already have for fully Bayesian inference but, you know, actually works 🙃 .

Regarding better support for `nn.Module` within Pyro's HMC: this would definitely be terrific to have. I imagine I'm not the only one that wants to take the parameters in a complex nested pre-existing `nn.Module` and do HMC over them.
Your workflow sounds reasonable and hopefully should unblock you for the time being, but one thing to check would be if this works with the JIT. Let us know if you hit any other blockers as well. I can take a look at better supporting NN modules within HMC next week which should also be straightforward.
@jacobrgardner another option might be to use [PyroModule](https://docs.pyro.ai/en/stable/nn.html#pyro.nn.module.PyroModule), which is the replacement for `pyro.random_module()`:
```py
# Define a torch.nn.Module.
model = ...

# Now make it Bayesian.
from pyro.nn.module import PyroSample, to_pyro_module_
to_pyro_module_(model)
model.covar_module.base_kernel.raw_lengthscale = PyroSample(Normal(0, 1).expand([1, 1]))
model.covar_module.raw_outputscale =  PyroSample(Normal(0, 1))
model.likelihood.noise_covar.raw_noise = PyroSample(Normal(0, 1).expand([1]))
model.mean_module.constant = PyroSample(Normal(0, 1))
```
@fritzo I initially tried that but couldn't get it working. For example:

```python
import math
import torch
import gpytorch
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.kernels.rbf_kernel import postprocess_rbf
from gpytorch.means import ConstantMean
from torch.nn.functional import softplus

from gpytorch.functions import RBFCovariance

import pyro
from pyro.distributions import Normal
from pyro.infer.mcmc import NUTS, HMC, MCMC


class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

x = torch.linspace(0, 1, 10).double()
y = torch.sin(x * (2 * math.pi)).double() + torch.randn(x.size()).double() * 0.1

likelihood = gpytorch.likelihoods.GaussianLikelihood()
likelihood.double()

model = ExactGPModel(x, y, likelihood)
model.double()

from pyro.nn.module import PyroSample, to_pyro_module_
to_pyro_module_(model)

def pyro_model6(x, y):
    model.covar_module.base_kernel.raw_lengthscale = PyroSample(Normal(0, 1).expand([1, 1]))
    model.covar_module.raw_outputscale =  PyroSample(Normal(0, 1))
    model.likelihood.noise_covar.raw_noise = PyroSample(Normal(0, 1).expand([1]))
    model.mean_module.constant = PyroSample(Normal(0, 1))
    
    output = model.likelihood(model(x))
    pyro.sample("obs", output, obs=y)
```

Seems to not work as expected, possibly because of something we do in `gpytorch.Module`. If I run:
```python
tr = trace(pyro_model6).get_trace(x, y)
display(tr.nodes)
```
The only nodes that come back are `_INPUT`, `obs` and `_RETURN`. I think at some point I may have gotten slightly further than this, but maybe the potential gradients were still also wrong?
@jacobrgardner thanks for explaining.  I wouldn't be surprised if there is some conflict between `PyroModule` and `gpytorch.Module`, especially as they have similar goals.  If @neerajprad solution works then great :+1:.  If you ever make further attempts at using `PyroModule`, don't hesitate to ping me on slack or on PRs.

Resolving patch:
diff --git a/gpytorch/mlls/exact_marginal_log_likelihood.py b/gpytorch/mlls/exact_marginal_log_likelihood.py
--- a/gpytorch/mlls/exact_marginal_log_likelihood.py
+++ b/gpytorch/mlls/exact_marginal_log_likelihood.py
@@ -65,10 +65,3 @@ def forward(self, function_dist, target, *params):
         # Scale by the amount of data we have
         num_data = function_dist.event_shape.numel()
         return res.div_(num_data)
-
-    def pyro_factor(self, output, target, *params):
-        import pyro
-
-        mll = target.size(-1) * self(output, target, *params)
-        pyro.factor("gp_mll", mll)
-        return mll
diff --git a/gpytorch/mlls/marginal_log_likelihood.py b/gpytorch/mlls/marginal_log_likelihood.py
--- a/gpytorch/mlls/marginal_log_likelihood.py
+++ b/gpytorch/mlls/marginal_log_likelihood.py
@@ -46,9 +46,3 @@ def forward(self, output, target, **kwargs):
         :param dict kwargs: Additional arguments to pass to the likelihood's :attr:`forward` function.
         """
         raise NotImplementedError
-
-    def pyro_factor(self, output, target):
-        """
-        As forward, but register the MLL with pyro using the pyro.factor primitive.
-        """
-        raise NotImplementedError
diff --git a/gpytorch/module.py b/gpytorch/module.py
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -1,7 +1,9 @@
 #!/usr/bin/env python3
 
+import copy
 import inspect
 import itertools
+import operator
 from collections import OrderedDict
 
 import torch
@@ -81,7 +83,6 @@ def initialize(self, **kwargs):
 
         Value can take the form of a tensor, a float, or an int
         """
-
         for name, val in kwargs.items():
             if isinstance(val, int):
                 val = float(val)
@@ -355,6 +356,21 @@ def sample_from_prior(self, prior_name):
             raise RuntimeError("Must provide inverse transform to be able to sample from prior.")
         setting_closure(self, prior.sample())
 
+    def to_pyro_random_module(self):
+        pyro_random_module_cls = type("_Pyro" + self.__class__.__name__, (RandomModuleMixin, self.__class__), {})
+        if not isinstance(self, pyro_random_module_cls):
+            new_module = copy.deepcopy(self)
+            new_module.__class__ = pyro_random_module_cls  # hack
+        else:
+            # Unclear if this branch would ever get used in practice, but it semantically makes sense to have.
+            new_module = copy.deepcopy(self)
+
+        for mname, child in new_module.named_children():
+            if isinstance(child, Module):
+                setattr(new_module, mname, child.to_pyro_random_module())
+
+        return new_module
+
     def pyro_sample_from_prior(self):
         """
         For each parameter in this Module and submodule that have defined priors, sample a value for that parameter
@@ -363,7 +379,8 @@ def pyro_sample_from_prior(self):
         This method can be used in a Pyro model to conveniently define pyro sample sites for all
         parameters of the model that have GPyTorch priors registered to them.
         """
-        return _pyro_sample_from_prior(module=self, memo=None, prefix="")
+        new_module = self.to_pyro_random_module()
+        return _pyro_sample_from_prior(module=new_module, memo=None, prefix="")
 
     def local_load_samples(self, samples_dict, memo, prefix):
         """
@@ -455,6 +472,7 @@ def _pyro_sample_from_prior(module, memo=None, prefix=""):
         import pyro
     except ImportError:
         raise RuntimeError("Cannot call pyro_sample_from_prior without pyro installed!")
+
     if memo is None:
         memo = set()
     if hasattr(module, "_priors"):
@@ -474,6 +492,8 @@ def _pyro_sample_from_prior(module, memo=None, prefix=""):
         submodule_prefix = prefix + ("." if prefix else "") + mname
         _pyro_sample_from_prior(module=module_, memo=memo, prefix=submodule_prefix)
 
+    return module
+
 
 def _pyro_load_from_samples(module, samples_dict, memo=None, prefix=""):
     if memo is None:
@@ -530,3 +550,43 @@ def _extract_named_constraints(module, memo=None, prefix=""):
         submodule_prefix = prefix + ("." if prefix else "") + mname
         for name, constraint in _extract_named_constraints(module_, memo=memo, prefix=submodule_prefix):
             yield name, constraint
+
+
+class RandomModuleMixin(object):
+    def initialize(self, **kwargs):
+        """
+        Set a value for a parameter
+
+        kwargs: (param_name, value) - parameter to initialize.
+        Can also initialize recursively by passing in the full name of a
+        parameter. For example if model has attribute model.likelihood,
+        we can initialize the noise with either
+        `model.initialize(**{'likelihood.noise': 0.1})`
+        or
+        `model.likelihood.initialize(noise=0.1)`.
+        The former method would allow users to more easily store the
+        initialization values as one object.
+
+        Value can take the form of a tensor, a float, or an int
+        """
+        for name, value in kwargs.items():
+            if not torch.is_tensor(value):
+                raise RuntimeError("Initialize in RandomModules can only be done with tensor values.")
+
+            names = name.rsplit(".")
+            if len(names) > 1:
+                mod_name, param_name = names
+                mod = operator.attrgetter(mod_name)(self)
+            else:
+                mod, param_name = self, name
+
+            old_param = getattr(mod, param_name)
+            is_property = hasattr(type(self), name) and isinstance(getattr(type(self), name), property)
+            if not isinstance(old_param, torch.nn.Parameter) or is_property:
+                # Presumably we're calling a getter that will call initialize again on the actual parameter.
+                setattr(mod, param_name, value.expand(old_param.shape))
+            else:
+                delattr(mod, param_name)
+                setattr(mod, param_name, value.expand(old_param.shape))
+
+        return self


Test patch:
diff --git a/test/examples/test_pyro_integration.py b/test/examples/test_pyro_integration.py
--- a/test/examples/test_pyro_integration.py
+++ b/test/examples/test_pyro_integration.py
@@ -116,6 +116,17 @@ def forward(self, x):
             res = gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
             return res
 
+    class ExactGPModel(gpytorch.models.ExactGP):
+        def __init__(self, train_x, train_y, likelihood):
+            super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
+            self.mean_module = gpytorch.means.ConstantMean()
+            self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+
+        def forward(self, x):
+            mean_x = self.mean_module(x)
+            covar_x = self.covar_module(x)
+            return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
     class LowLevelInterfaceClusterMultitaskGPModel(gpytorch.models.ApproximateGP):
         def __init__(self, train_x, train_y, num_functions=2):
             # Define all the variational stuff
@@ -301,6 +312,54 @@ def test_low_level_interface(self):
             self.assertGreater(cluster_probs[2, cluster_2_idx].squeeze().item(), 0.9)
             self.assertGreater(cluster_probs[3, cluster_1_idx].squeeze().item(), 0.9)
 
+        def test_nuts_derivatives(self):
+            x = torch.linspace(0, 1, 10).double()
+            y = torch.sin(x * (2 * pi)).double() + torch.randn(x.size()).double() * 0.1
+
+            likelihood = gpytorch.likelihoods.GaussianLikelihood()
+            likelihood.double()
+
+            model = ExactGPModel(x, y, likelihood)
+            model.double()
+
+            model.covar_module.base_kernel.register_prior(
+                "lengthscale_prior", gpytorch.priors.LogNormalPrior(0.1, 1.), "lengthscale"
+            )
+            model.covar_module.register_prior(
+                "raw_outputscale_prior", gpytorch.priors.NormalPrior(0., 1.), "raw_outputscale"
+            )
+            model.likelihood.register_prior("raw_noise_prior", gpytorch.priors.NormalPrior(0., 1.), "raw_noise")
+            model.mean_module.register_prior("constant_prior", gpytorch.priors.NormalPrior(0., 1.), "constant")
+
+            def pyro_model(x, y):
+                sampled_model = model.pyro_sample_from_prior()
+                output = sampled_model.likelihood(sampled_model(x))
+                pyro.sample("obs", output, obs=y)
+
+            nuts_kernel = pyro.infer.mcmc.NUTS(pyro_model)
+            nuts_kernel.setup(0, x, y)
+
+            for i in range(10):
+                # Sampling some random values for each of the parameters
+                sampled_model = model.pyro_sample_from_prior()
+                initial_params = {
+                    'covar_module.base_kernel.lengthscale_prior': sampled_model.covar_module.base_kernel.lengthscale,
+                    'covar_module.raw_outputscale_prior': sampled_model.covar_module.raw_outputscale,
+                    'mean_module.constant_prior': sampled_model.mean_module.constant,
+                    'likelihood.raw_noise_prior': sampled_model.likelihood.raw_noise}
+
+                # check gradient of the potential fn w.r.t. the parameters at these specific values.
+                d = 1e-4
+                grads, v = pyro.ops.integrator.potential_grad(nuts_kernel.potential_fn, initial_params)
+                for param in initial_params:
+                    new_params = initial_params.copy()
+                    new_params[param] = new_params[param] + d
+                    v2 = nuts_kernel.potential_fn(new_params).item()
+                    numerical = (v2 - v) / d
+
+                    abs_delta = (torch.abs(grads[param] - numerical) / torch.abs(numerical)).item()
+                    # Pyro's gradient should differ by less than 5% from numerical est.
+                    self.assertLess(abs_delta, 0.05)
 
 except ImportError:
     pass
diff --git a/test/examples/test_simple_gp_regression.py b/test/examples/test_simple_gp_regression.py
--- a/test/examples/test_simple_gp_regression.py
+++ b/test/examples/test_simple_gp_regression.py
@@ -459,16 +459,15 @@ def test_pyro_sampling(self):
         gp_model.covar_module.register_prior("outputscale_prior", UniformPrior(1, 2), "outputscale")
         likelihood.register_prior("noise_prior", UniformPrior(0.05, 0.3), "noise")
 
-        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
-
         def pyro_model(x, y):
-            gp_model.pyro_sample_from_prior()
-            output = gp_model(x)
-            mll.pyro_factor(output, y)
+            with gpytorch.settings.fast_computations(False, False, False):
+                sampled_model = gp_model.pyro_sample_from_prior()
+                output = sampled_model.likelihood(sampled_model(x))
+                pyro.sample("obs", output, obs=y)
             return y
 
         nuts_kernel = NUTS(pyro_model, adapt_step_size=True)
-        mcmc_run = MCMC(nuts_kernel, num_samples=3, warmup_steps=20)
+        mcmc_run = MCMC(nuts_kernel, num_samples=3, warmup_steps=20, disable_progbar=True)
         mcmc_run.run(train_x, train_y)
 
         gp_model.pyro_load_from_samples(mcmc_run.get_samples())