Implement parallel enumeration over discrete sample sites #776

fritzo · 2018-02-15T02:58:03Z

See Design Doc | Closes #742 | Replaces #227

This implements parallel enumeration for discrete sample sites. The new notation is

- SVI(model, guide, optim, "ELBO", enum_discrete=True)
+ SVI(model, config_enumerate(guide), optim, "ELBO")

This allows finer-grained control over each discrete sample site, e.g. we can freely mix monte carlo sampling, sequential enumeration, and parallel enumeration:

def guide():
    pyro.sample("x", Categorical(p), infer={"enumerate": "sequential"})
    pyro.sample("y", Categorical(p), infer={"enumerate": "parallel"})
    pyro.sample("z", Categorical(p))  # monte carlo

Additionally we can use the @config_enumerate decorator to annotate an entire guide:

@config_enumerate
def guide():
    pyro.sample("x", Categorical(p))  # sequential
    pyro.sample("y", Categorical(p), infer={"enumerate": "parallel"})
    pyro.sample("z", Categorical(p), infer={"enumerate": None})  # monte carlo

We can also set parallel enumeration by default

@config_enumerate(default="parallel")
def guide():
    ...

Tasks

Add tests for max_iarange_nesting
Fix batch_log_pdf.shape errors
Add tests for nested enumeration (mixing sequential with parallel)
Add tests for correctness of gradients

fritzo · 2018-02-21T21:55:48Z

@eb8680 Could you take a look at the new interface using enumerate_discrete(guide) and your Poutine rather than the enum_discrete=True kwarg? (there are still some test failures, but the interface is mostly complete).

eb8680 · 2018-02-21T23:26:44Z

Could you take a look at the new interface using enumerate_discrete(guide) and your Poutine rather than the enum_discrete=True kwarg?

It mostly looks fine, but maybe we should change the name of enumerate_discrete to config_enumerate or something similar to make it clear that it does not modify the semantics of the guide on its own (i.e. calling enumerate_discrete(guide)(...) with no enclosing EnumerateMessenger context will sample the latent variables rather than enumerating).

fritzo · 2018-02-23T02:54:26Z

tests/infer/test_enum.py

    with xfail_if_not_implemented():
        inference.step(data)


+@pytest.mark.parametrize("enum_discrete", [None, "sequential", "parallel"])
+@pytest.mark.parametrize("trace_graph", [False, True], ids=["dense", "flat"])
+def test_bern_elbo_gradient(enum_discrete, trace_graph):


This and the following test use analytic grad(kl_divergence(-, -), -) to test correctness of ELBO gradients.

fritzo · 2018-02-23T02:59:15Z

pyro/infer/tracegraph_elbo.py

@@ -190,9 +190,6 @@ def _get_traces(self, model, guide, *args, **kwargs):
        """

        for i in range(self.num_particles):
-            if self.enum_discrete:


Enumeration is now controlled by hints in the site['infer'] dict. If an inference algorithm does not implement enumeration, it can safely ignore those hints and sample rather than enumerate.

neerajprad · 2018-02-24T00:06:31Z

pyro/infer/enum.py

-        q_fn = poutine.queue(fn, queue=queue)
-        full_trace = poutine.trace(
-            q_fn, graph_type=graph_type).get_trace(*args, **kwargs)
+        q_fn = poutine.queue(fn, queue=queue, escape_fn=_iter_discrete_escape)


Should this be outside the while loop replacing q_fn above?

That looks right to me. Done.

neerajprad · 2018-02-24T00:12:09Z

pyro/distributions/util.py

@@ -90,40 +90,58 @@ def sum_rightmost(value, dim):
    """
    Sum out ``dim`` many rightmost dimensions of a given tensor.

+    If ``dim`` is 0, no dimensions are summed out.
+    If ``dim`` is ``float('inf')``, then all dimensions are summed out.
+    If ``dim`` is 1, the leftmost 1 dimension is summed out.


leftmost -> rightmost and reverse below.

neerajprad · 2018-02-24T00:27:04Z

pyro/infer/enum.py



-def site_is_discrete(name, site):
-    return getattr(site["fn"], "enumerable", False)
+def _iter_discrete_filter(name, msg):


Can we remove the name from the function args?

I can't remove this arg without changes to Trace.compute_batch_log_pdf(). Let's clean this up later when we refactor the Trace.compute_() methods.

neerajprad · 2018-02-24T00:27:35Z

pyro/poutine/enumerate_poutine.py

+from .poutine import Messenger, Poutine
+
+
+def _iter_discrete_filter(name, msg):


Same here. Can we remove name?

fritzo · 2018-02-24T05:07:42Z

Woo hoo, tests finally pass!

martinjankowiak · 2018-02-24T19:01:51Z

pyro/infer/elbo.py

@@ -25,9 +29,9 @@ class ELBO(object):

    def __init__(self,
                 num_particles=1,
-                 enum_discrete=False):
+                 max_iarange_nesting=float('inf')):


is it weird to mix floats and ints?

It seems pretty safe to me. We're taking care to consume this number only through comparisons (in functions like sum_rightmost() and check_sites() in #806 ) and we only read its value if it is less than some other finite number. This float('inf') solution seems cleaner to me than the alternatives like INT_MAX or -1 for which we would need logic that is even less readable.

martinjankowiak · 2018-02-24T19:04:44Z

pyro/infer/enum.py

+
+def config_enumerate(guide=None, default="sequential"):
+    """
+    Configures each enumerable site a guide to enumerate with given method,


should the docstring mention that it doesn't override?

martinjankowiak · 2018-02-24T19:05:12Z

pyro/infer/trace_elbo.py

-               weight.dim() > 0 and \
-               weight.size(0) > 1
+            # iterate over a bag of traces, one trace per particle
+            for scale, guide_trace in iter_discrete_traces("flat", self.max_iarange_nesting, guide, *args, **kwargs):


so beautiful

martinjankowiak · 2018-02-24T19:09:43Z

pyro/poutine/enumerate_poutine.py

+from .poutine import Messenger, Poutine
+
+
+def _iter_discrete_filter(msg):


this is defined in two places

this = _iter_discrete_filter

They differ, but I'll rename them to make that clear...

martinjankowiak · 2018-02-24T19:13:57Z

pyro/poutine/enumerate_poutine.py

+
+            # Ensure enumeration happens at an available tensor dimension.
+            event_dim = len(msg["fn"].event_shape)
+            actual_dim = value.dim() - event_dim - 1


can we have one or two comments here i find this confusing

actual_dim = value.dim() - event_dim - 1

Ok, I've simplified and added some comments.

martinjankowiak · 2018-02-24T19:23:05Z

tests/infer/test_enum.py

@@ -49,6 +50,7 @@ def model():
 def test_iter_discrete_traces_vector(graph_type):


what's the deal with test_iter_discrete_traces_vector?

Fixed and removed the @xfail

martinjankowiak

lgtm!

great work!!!

fritzo · 2018-02-25T18:33:53Z

Finally ready to merge!

fritzo added 3 commits February 14, 2018 18:17

Sketch EnumeratePoutine

369ecb1

Merge branch 'dev' into enumerate-parallel

403708b

Fix dimension logic in EnumerateMessenger

95957e6

fritzo added the WIP label Feb 15, 2018

fritzo requested review from eb8680 and martinjankowiak February 15, 2018 02:58

fritzo added 2 commits February 15, 2018 09:16

Add more test examples

ca6da58

Refactor ELBO

c3c0c80

fritzo force-pushed the enumerate-parallel branch from d352788 to c3c0c80 Compare February 15, 2018 17:55

fritzo added 2 commits February 20, 2018 09:52

Merge branch 'dev' into enumerate-parallel

db45ca5

Attempt to get batch shapes correct for enum_discrete in trace_elbo

c56cd6a

fritzo mentioned this pull request Feb 20, 2018

Parallelize ELBO computation over num_particles #791

Closed

fritzo added 2 commits February 20, 2018 14:29

Merge branch 'dev' into enumerate-parallel

a1f2b14

Simplify Trace_ELBO

c39d44a

fritzo mentioned this pull request Feb 20, 2018

Infer config poutine #792

Merged

fritzo added 5 commits February 20, 2018 16:06

Drop special-case for enum_discrete in Trace_ELBO

1b60102

Merge branch 'dev' into enumerate-parallel

ebbe9d1

Replace enum_discrete kwarg with enumerate_discrete() function

49606f8

Completely elimitate enum_discrete kwarg

2a7e25d

Fix bugs in tests/infer/test_enum.py

0da83b9

fritzo added 3 commits February 21, 2018 16:32

Rename enumerate_discrete to config_enumerate

45140f4

Merge branch 'dev' into enumerate-parallel

53a44ee

Add analytic KL tests for parallel enumeration

37b6ac9

fritzo force-pushed the enumerate-parallel branch from 626c59b to 37b6ac9 Compare February 23, 2018 02:41

fritzo added awaiting review and removed WIP labels Feb 23, 2018

Add test for sum_rightmost()

5ff4c32

fritzo commented Feb 23, 2018

View reviewed changes

fritzo added 3 commits February 22, 2018 19:26

Skip slow tests on travis

7787050

Add another gradient test for enumeration

fe8820c

Add TODOs for more tests

2a52b80

fritzo mentioned this pull request Feb 23, 2018

Adopt strict batch shape semantics for distributions #806

Merged

7 tasks

fritzo added 2 commits February 23, 2018 13:47

Add variously-sized categoricals test

c61a1ba

Remove excruciatingly slow test

5e9e4ff

neerajprad reviewed Feb 24, 2018

View reviewed changes

fritzo added 2 commits February 23, 2018 16:49

Fix scalar error

1947ba0

Flake8

0d695db

neerajprad reviewed Feb 24, 2018

View reviewed changes

fritzo added 2 commits February 23, 2018 18:11

Remove name arg to _iter_discrete_filter

617af11

Updates per review

2ae92cf

martinjankowiak requested changes Feb 24, 2018

View reviewed changes

Updates per review

1e26198

martinjankowiak approved these changes Feb 25, 2018

View reviewed changes

eb8680 approved these changes Feb 25, 2018

View reviewed changes

fritzo mentioned this pull request Feb 25, 2018

Enumerate continuous RVs using Monte Carlo or quadrature #811

Closed

martinjankowiak merged commit 1bbfb48 into dev Feb 25, 2018

fritzo mentioned this pull request Feb 28, 2018

Fix enumerate + iarange #828

Merged

4 tasks

fritzo deleted the enumerate-parallel branch March 6, 2018 23:51

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement parallel enumeration over discrete sample sites #776

Implement parallel enumeration over discrete sample sites #776

fritzo commented Feb 15, 2018 •

edited

Loading

fritzo commented Feb 21, 2018

eb8680 commented Feb 21, 2018

fritzo Feb 23, 2018 •

edited

Loading

fritzo Feb 23, 2018 •

edited

Loading

neerajprad Feb 24, 2018

fritzo Feb 24, 2018

neerajprad Feb 24, 2018

neerajprad Feb 24, 2018

fritzo Feb 24, 2018

neerajprad Feb 24, 2018

fritzo Feb 24, 2018

fritzo commented Feb 24, 2018

martinjankowiak Feb 24, 2018

fritzo Feb 25, 2018

martinjankowiak Feb 24, 2018

fritzo Feb 25, 2018

martinjankowiak Feb 24, 2018

martinjankowiak Feb 24, 2018

martinjankowiak Feb 24, 2018

fritzo Feb 25, 2018

martinjankowiak Feb 24, 2018

fritzo Feb 25, 2018

martinjankowiak Feb 24, 2018

fritzo Feb 25, 2018

martinjankowiak left a comment

fritzo commented Feb 25, 2018

		from .poutine import Messenger, Poutine


		def _iter_discrete_filter(name, msg):

		from .poutine import Messenger, Poutine


		def _iter_discrete_filter(msg):

		@@ -49,6 +50,7 @@ def model():
		def test_iter_discrete_traces_vector(graph_type):

Implement parallel enumeration over discrete sample sites #776

Implement parallel enumeration over discrete sample sites #776

Conversation

fritzo commented Feb 15, 2018 • edited Loading

Tasks

fritzo commented Feb 21, 2018

eb8680 commented Feb 21, 2018

fritzo Feb 23, 2018 • edited Loading

Choose a reason for hiding this comment

fritzo Feb 23, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fritzo commented Feb 24, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

martinjankowiak left a comment

Choose a reason for hiding this comment

fritzo commented Feb 25, 2018

fritzo commented Feb 15, 2018 •

edited

Loading

fritzo Feb 23, 2018 •

edited

Loading

fritzo Feb 23, 2018 •

edited

Loading