Repository: pyro-ppl/pyro

Issue:
Feature request: display ADVI params
Hi,
When using ADVI guide, it is possible to print nicely a sample of the model parameters, with e.g.
print(advi.guide(y))
Pyro's param store contains advi_loc and advi_scale/advi_scale_tril, but it is not clear which parameter refers to which index of these vectors/arrays and, from what I understand, these are unrestricted params, not my models restricted params I care about.
So, I would like to request adding a functionality for a nice display of the means of location and scale of the model parameters.

Slawek


Discussion:
I think the easiest interface would be to expose quantiles, e.g. `.median()`. Some of these parameters are transformed, and it's not obvious how to interpret e.g. the variance of a normal distribution that is then transformed via sigmoid transoform to lie in `[0,1]`. I'll expose `.median()` first.

Resolving patch:
diff --git a/pyro/infer/advi.py b/pyro/infer/advi.py
--- a/pyro/infer/advi.py
+++ b/pyro/infer/advi.py
@@ -77,6 +77,22 @@ def sample_latent(self, *args, **kwargs):
         """
         raise NotImplementedError
 
+    def _unpack_latent(self, latent):
+        """
+        Unpacks a packed latent tensor, iterating over tuples of the form::
+
+            (site, unconstrained_value)
+        """
+        pos = 0
+        for name, site in self.prototype_trace.nodes.items():
+            if site["type"] == "sample" and not site["is_observed"]:
+                unconstrained_shape = self._unconstrained_shapes[name]
+                size = _product(unconstrained_shape)
+                unconstrained_value = latent[pos:pos + size].view(unconstrained_shape)
+                yield site, unconstrained_value
+                pos += size
+        assert pos == len(latent)
+
     def guide(self, *args, **kwargs):
         """
         An automatic guide with the same ``*args, **kwargs`` as the base ``model``.
@@ -88,7 +104,6 @@ def guide(self, *args, **kwargs):
         if self.prototype_trace is None:
             self._setup_prototype(*args, **kwargs)
 
-        result = {}
         latent = self.sample_latent(*args, **kwargs)
 
         # create all iaranges
@@ -96,25 +111,19 @@ def guide(self, *args, **kwargs):
                     for frame in sorted(self._iaranges.values())}
 
         # unpack latent samples
-        pos = 0
-        for name, site in self.prototype_trace.nodes.items():
-            if site["type"] == "sample" and not site["is_observed"]:
-                unconstrained_shape = self._unconstrained_shapes[name]
-                size = _product(unconstrained_shape)
-                unconstrained_value = latent[pos:pos + size].view(unconstrained_shape)
-                pos += size
-                transform = biject_to(site["fn"].support)
-                value = transform(unconstrained_value)
-                log_density = transform.inv.log_abs_det_jacobian(value, unconstrained_value)
-                log_density = sum_rightmost(log_density, log_density.dim() - value.dim() + site["fn"].event_dim)
-                delta_dist = dist.Delta(value, log_density=log_density, event_dim=site["fn"].event_dim)
-
-                with ExitStack() as stack:
-                    for frame in self._cond_indep_stacks[name]:
-                        stack.enter_context(iaranges[frame.name])
-                    result[name] = pyro.sample(name, delta_dist)
-
-        assert pos == len(latent)
+        result = {}
+        for site, unconstrained_value in self._unpack_latent(latent):
+            name = site["name"]
+            transform = biject_to(site["fn"].support)
+            value = transform(unconstrained_value)
+            log_density = transform.inv.log_abs_det_jacobian(value, unconstrained_value)
+            log_density = sum_rightmost(log_density, log_density.dim() - value.dim() + site["fn"].event_dim)
+            delta_dist = dist.Delta(value, log_density=log_density, event_dim=site["fn"].event_dim)
+
+            with ExitStack() as stack:
+                for frame in self._cond_indep_stacks[name]:
+                    stack.enter_context(iaranges[frame.name])
+                result[name] = pyro.sample(name, delta_dist)
 
         return result
 
@@ -161,6 +170,39 @@ def sample_latent(self, *args, **kwargs):
                                 constraint=constraints.lower_cholesky)
         return pyro.sample("_advi_latent", dist.MultivariateNormal(loc, scale_tril=scale_tril))
 
+    def median(self, *args, **kwargs):
+        """
+        Returns the posterior median value of each latent variable.
+
+        :return: A dictionary mapping sample site name to median tensor.
+        :rtype: dict
+        """
+        latent = pyro.param("advi_loc")
+        return {site["name"]: biject_to(site["fn"].support)(unconstrained_value)
+                for site, unconstrained_value in self._unpack_latent(latent)}
+
+    def quantiles(self, quantiles, *args, **kwargs):
+        """
+        Returns posterior quantiles each latent variable. Example::
+
+            print(advi.quantiles([0.05, 0.5, 0.95]))
+
+        :param quantiles: A list of requested quantiles between 0 and 1.
+        :type quantiles: torch.Tensor or list
+        :return: A dictionary mapping sample site name to a list of quantile values.
+        :rtype: dict
+        """
+        loc = pyro.param("advi_loc")
+        scale = pyro.param("advi_scale_tril").diag()
+        quantiles = loc.new_tensor(quantiles).unsqueeze(-1)
+        latents = dist.Normal(loc, scale).icdf(quantiles)
+
+        result = {}
+        for latent in latents:
+            for site, unconstrained_value in self._unpack_latent(latent):
+                result.setdefault(site["name"], []).append(biject_to(site["fn"].support)(unconstrained_value))
+        return result
+
 
 class ADVIDiagonalNormal(ADVI):
     """
@@ -190,3 +232,36 @@ def sample_latent(self, *args, **kwargs):
         scale = pyro.param("advi_scale", torch.ones(self.latent_dim),
                            constraint=constraints.positive)
         return pyro.sample("_advi_latent", dist.Normal(loc, scale).reshape(extra_event_dims=1))
+
+    def median(self, *args, **kwargs):
+        """
+        Returns the posterior median value of each latent variable.
+
+        :return: A dictionary mapping sample site name to median tensor.
+        :rtype: dict
+        """
+        latent = pyro.param("advi_loc")
+        return {site["name"]: biject_to(site["fn"].support)(unconstrained_value)
+                for site, unconstrained_value in self._unpack_latent(latent)}
+
+    def quantiles(self, quantiles, *args, **kwargs):
+        """
+        Returns posterior quantiles each latent variable. Example::
+
+            print(advi.quantiles([0.05, 0.5, 0.95]))
+
+        :param quantiles: A list of requested quantiles between 0 and 1.
+        :type quantiles: torch.Tensor or list
+        :return: A dictionary mapping sample site name to a list of quantile values.
+        :rtype: dict
+        """
+        loc = pyro.param("advi_loc")
+        scale = pyro.param("advi_scale")
+        quantiles = loc.new_tensor(quantiles).unsqueeze(-1)
+        latents = dist.Normal(loc, scale).icdf(quantiles)
+
+        result = {}
+        for latent in latents:
+            for site, unconstrained_value in self._unpack_latent(latent):
+                result.setdefault(site["name"], []).append(biject_to(site["fn"].support)(unconstrained_value))
+        return result


Test patch:
diff --git a/tests/infer/test_advi.py b/tests/infer/test_advi.py
--- a/tests/infer/test_advi.py
+++ b/tests/infer/test_advi.py
@@ -7,7 +7,9 @@
 import pyro
 import pyro.distributions as dist
 import pyro.poutine as poutine
-from pyro.infer import ELBO, ADVIDiagonalNormal, ADVIMultivariateNormal
+from pyro.infer import ELBO, SVI, ADVIDiagonalNormal, ADVIMultivariateNormal
+from pyro.optim import Adam
+from tests.common import assert_equal
 
 
 @pytest.mark.parametrize("advi_class", [ADVIMultivariateNormal, ADVIDiagonalNormal])
@@ -44,3 +46,56 @@ def model():
     elbo = ELBO.make(trace_graph=trace_graph, enum_discrete=enum_discrete)
     loss = elbo.loss(advi.model, advi.guide)
     assert np.isfinite(loss), loss
+
+
+@pytest.mark.parametrize("advi_class", [ADVIMultivariateNormal, ADVIDiagonalNormal])
+def test_median(advi_class):
+
+    def model():
+        pyro.sample("x", dist.Normal(0.0, 1.0))
+        pyro.sample("y", dist.LogNormal(0.0, 1.0))
+        pyro.sample("z", dist.Beta(2.0, 2.0))
+
+    advi = advi_class(model)
+    infer = SVI(advi.model, advi.guide, Adam({'lr': 0.01}), 'ELBO')
+    for _ in range(100):
+        infer.step()
+
+    median = advi.median()
+    assert_equal(median["x"], torch.tensor(0.0), prec=0.1)
+    assert_equal(median["y"], torch.tensor(1.0), prec=0.1)
+    assert_equal(median["z"], torch.tensor(0.5), prec=0.1)
+
+
+@pytest.mark.parametrize("advi_class", [ADVIMultivariateNormal, ADVIDiagonalNormal])
+def test_quantiles(advi_class):
+
+    def model():
+        pyro.sample("x", dist.Normal(0.0, 1.0))
+        pyro.sample("y", dist.LogNormal(0.0, 1.0))
+        pyro.sample("z", dist.Beta(2.0, 2.0))
+
+    advi = advi_class(model)
+    infer = SVI(advi.model, advi.guide, Adam({'lr': 0.01}), 'ELBO')
+    for _ in range(100):
+        infer.step()
+
+    quantiles = advi.quantiles([0.1, 0.5, 0.9])
+    median = advi.median()
+    for name in ["x", "y", "z"]:
+        assert_equal(median[name], quantiles[name][1])
+
+    assert torch.tensor(-3.0) < quantiles["x"][0]
+    assert quantiles["x"][0] + 1.0 < quantiles["x"][1]
+    assert quantiles["x"][1] + 1.0 < quantiles["x"][2]
+    assert quantiles["x"][2] < 3.0
+
+    assert torch.tensor(0.01) < quantiles["y"][0]
+    assert quantiles["y"][0] * 2.0 < quantiles["y"][1]
+    assert quantiles["y"][1] * 2.0 < quantiles["y"][2]
+    assert quantiles["y"][2] < 100.0
+
+    assert torch.tensor(0.01) < quantiles["z"][0]
+    assert quantiles["z"][0] + 0.1 < quantiles["z"][1]
+    assert quantiles["z"][1] + 0.1 < quantiles["z"][2]
+    assert quantiles["z"][2] < 0.99