Repository: pyro-ppl/pyro

Issue:
Make autoguides more usable
This aims to make our most common autoguides easier to use: `AutoDelta`, `AutoDiagonalNormal`, `AutoMultivariateNormal`, `AutoLowRankMultivariateNormal`.

Feel free to add to this list (anyone :smile:) if you have usability requests.

I have added this to the 1.0 milestone because it involves interface changes that will enable later refactoring (e.g. #2078) while preserving backwards compatibility.

## Tasks
- [x] #2119 document use of custom `init_loc_fn`, and add examples in tutorials
- [x] #2119 add error checking logic in `init_loc_fn`
- [x] #2119 add easier `init_scale` args for `Auto*Normal` guides
- [x] #2126 make autoguides into `nn.Module`s so they are easier to `torch.load()`/`.save()`
    (this is an alternative to making the param store an `nn.Module` #2078 )
- [x] allow non-python-identifiers as site names in autoguides, ~~which was broken by #2126~~
- ~~make it easier to set optim configs per-guide-param (e.g. I often lower the learning rate `lr *= 0.1` for the `auto_cov_factor` param of `AutoLowRankMultivariateNormal`).~~
- [x] #2127 Stabilize parameterization of `AutoLowRankMultivariateNormal`
- [x] #2129 Support pickling of easyguides
- [x] #2137 Make `AutoGuideList` a bonafide `nn.ModuleList`
- [x] #2140 make easyguides into `nn.Modules`
- [x] #2148 Fix saving of `AutoGuides` in the param store
- [ ] #2156 Change default `init_scale` from 1.0 to 0.1 (@fritzo often uses this in practice)


Discussion:
@martinjankowiak @stefanwebb let me know if you have any other suggestions!
@fritzo Could we have a method in PyroModule to get constrained values (without having to use `pyro.get_param_store()`)? Right now, we can observe unconstrained parameters through `.named_parameters()` method but it would be nice to have support for constrained parameters too. :)
@fehiepsi sure, feel free to add a `.get_named_pyro_params()` or similar.
@fritzo Just discover a bug with the dev branch :D 
```python
class Child(PyroModule):
    def __init__(self, x):
        super().__init__()
        self.a = PyroParam(torch.tensor(x))

class Family(PyroModule):
    def __init__(self):
        super().__init__()
        self.c = Child(1.0)

    def forward(self):
        return self.c.a

f.c = Child(3.)
assert f() == 3
```
@fehiepsi good point, I believe that's because the `__delattr__` logic currently only removes `PyroParam` and `nn.Param` from the param store (`__delattr__` is triggered when `__setattr__` overwrites something). Do you want to try modifying `__delattr__` to correctly handle `PyroModule` and `nn.Module`? That seems like a good first PR contributing to `PyroModule` :smile:

Resolving patch:
diff --git a/pyro/infer/autoguide/guides.py b/pyro/infer/autoguide/guides.py
--- a/pyro/infer/autoguide/guides.py
+++ b/pyro/infer/autoguide/guides.py
@@ -517,7 +517,7 @@ class AutoMultivariateNormal(AutoContinuous):
         svi = SVI(model, guide, ...)
 
     By default the mean vector is initialized by ``init_loc_fn()`` and the
-    Cholesky factor is initialized to the identity.
+    Cholesky factor is initialized to the identity times a small factor.
 
     :param callable model: A generative model.
     :param callable init_loc_fn: A per-site initialization function.
@@ -526,8 +526,7 @@ class AutoMultivariateNormal(AutoContinuous):
         (unconstrained transformed) latent variable.
     """
 
-    def __init__(self, model, init_loc_fn=init_to_median,
-                 init_scale=1.0):
+    def __init__(self, model, init_loc_fn=init_to_median, init_scale=0.1):
         if not isinstance(init_scale, float) or not (init_scale > 0):
             raise ValueError("Expected init_scale > 0. but got {}".format(init_scale))
         self._init_scale = init_scale
@@ -562,7 +561,7 @@ class AutoDiagonalNormal(AutoContinuous):
         svi = SVI(model, guide, ...)
 
     By default the mean vector is initialized to zero and the scale is
-    initialized to the identity.
+    initialized to the identity times a small factor.
 
     :param callable model: A generative model.
     :param callable init_loc_fn: A per-site initialization function.
@@ -571,8 +570,7 @@ class AutoDiagonalNormal(AutoContinuous):
         (unconstrained transformed) latent variable.
     """
 
-    def __init__(self, model, init_loc_fn=init_to_median,
-                 init_scale=1.0):
+    def __init__(self, model, init_loc_fn=init_to_median, init_scale=0.1):
         if not isinstance(init_scale, float) or not (init_scale > 0):
             raise ValueError("Expected init_scale > 0. but got {}".format(init_scale))
         self._init_scale = init_scale
@@ -607,22 +605,24 @@ class AutoLowRankMultivariateNormal(AutoContinuous):
         guide = AutoLowRankMultivariateNormal(model, rank=10)
         svi = SVI(model, guide, ...)
 
-    By default the ``cov_diag`` is initialized to 1/2 and the ``cov_factor`` is
-    intialized randomly such that ``cov_factor.matmul(cov_factor.t())`` is half the
-    identity matrix.
+    By default the ``cov_diag`` is initialized to a small constant and the
+    ``cov_factor`` is initialized randomly such that on average
+    ``cov_factor.matmul(cov_factor.t())`` has the same scale as ``cov_diag``.
 
     :param callable model: A generative model.
-    :param int rank: The rank of the low-rank part of the covariance matrix.
+    :param rank: The rank of the low-rank part of the covariance matrix.
+        Defaults to approximately ``sqrt(latent dim)``.
+    :type rank: int or None
     :param callable init_loc_fn: A per-site initialization function.
         See :ref:`autoguide-initialization` section for available functions.
     :param float init_scale: Approximate initial scale for the standard
         deviation of each (unconstrained transformed) latent variable.
     """
 
-    def __init__(self, model, init_loc_fn=init_to_median, init_scale=1.0, rank=1):
+    def __init__(self, model, init_loc_fn=init_to_median, init_scale=0.1, rank=None):
         if not isinstance(init_scale, float) or not (init_scale > 0):
             raise ValueError("Expected init_scale > 0. but got {}".format(init_scale))
-        if not isinstance(rank, int) or not rank > 0:
+        if not (rank is None or isinstance(rank, int) and rank > 0):
             raise ValueError("Expected rank > 0 but got {}".format(rank))
         self._init_scale = init_scale
         self.rank = rank
@@ -632,6 +632,8 @@ def _setup_prototype(self, *args, **kwargs):
         super()._setup_prototype(*args, **kwargs)
         # Initialize guide params
         self.loc = nn.Parameter(self._init_loc())
+        if self.rank is None:
+            self.rank = int(round(self.latent_dim ** 0.5))
         self.scale = PyroParam(
             self.loc.new_full((self.latent_dim,), 0.5 ** 0.5 * self._init_scale),
             constraint=constraints.positive)


Test patch:
diff --git a/tests/infer/test_autoguide.py b/tests/infer/test_autoguide.py
--- a/tests/infer/test_autoguide.py
+++ b/tests/infer/test_autoguide.py
@@ -216,8 +216,11 @@ def model():
         pyro.sample("z", dist.Beta(2.0, 2.0))
 
     guide = auto_class(model)
-    infer = SVI(model, guide, Adam({'lr': 0.005}), Elbo(strict_enumeration_warning=False))
-    for _ in range(800):
+    optim = Adam({'lr': 0.05, 'betas': (0.8, 0.99)})
+    elbo = Elbo(strict_enumeration_warning=False,
+                num_particles=100, vectorize_particles=True)
+    infer = SVI(model, guide, optim, elbo)
+    for _ in range(100):
         infer.step()
 
     if auto_class is AutoLaplaceApproximation:
@@ -303,7 +306,10 @@ def model():
         pyro.sample("z", dist.Beta(2.0, 2.0))
 
     guide = auto_class(model)
-    infer = SVI(model, guide, Adam({'lr': 0.01}), Elbo(strict_enumeration_warning=False))
+    optim = Adam({'lr': 0.05, 'betas': (0.8, 0.99)})
+    elbo = Elbo(strict_enumeration_warning=False,
+                num_particles=100, vectorize_particles=True)
+    infer = SVI(model, guide, optim, elbo)
     for _ in range(100):
         infer.step()