Repository: pyro-ppl/pyro
Pull request: Support custom batching in iarange/irange
Closes #169

This aims to implement custom batching strategies in a `poutine.replay`-compatible way, as discussed in #239 .

It currently does not work, since observations override replaying. @eb8680 any ideas how to fix? 

Problem statement:
more control over subsampling in map_data
currently, mini-batch indices for map_data subsampling are chosen by pyro. they are also chosen in a stateless way with the consequence that there is non-zero probability that certain data points will never be seen even no matter how many epochs you do. this is bad, especially for large data sets. the simplest solution to this problem would be to give the user the ability to specify `batch_indices` directly in `map_data` (a solution that would presumably be psychologically comforting to users who are used to doing this anyway). another solution, which could coexist with the first solution, would be for the backend to implement more subsampling strategies besides the 'naive' one that we currently have.

Hints from issue discussion:
If the user wants to directly specify batch_indices, they can just not use `map_data` and write what they want directly.  Maybe we can add an option to the default `Poutine` to sample without replacement in `map_data`, though.

We could also expose the randomness in `map_data` as a Pyro variable, to enable fancy active subsampling strategies.
unless i'm missing something in that case scale factors won't be propagated and any other benefits from map_data (like independency information used for rao-blackwellization) fall by the wayside
Hmm... I guess not, although this seems like a potential source of nasty statistical bugs, especially in the nested case, or when there's a sequence of variables following a `map_data` that depend on all the variables inside.
yeah, that's certainly true. but my general design philosophy remains:

- try to make the default/highlighted interfaces relatively robust to user error
- but in those cases where finer control can be of a lot of practical importance, allow the advanced user finer knobs even if this opens the door to bugs and math errors. potential pitfalls can also be highlighted with friendly warnings as in #164 
i agree with @martinjankowiak. it's pretty common, eg, to want to iterate over the batches of data.

we should also consider doing this extension in such a way that it plays nice with dynamic data loaders for when data is too big to fit in memory....

Commit message: Attempt to implement custom iarange/irange subsampling

Patch to review:
diff --git a/pyro/__init__.py b/pyro/__init__.py
index 293e1617ba..92c0b16b75 100644
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -117,7 +117,7 @@ def observe(name, fn, obs, *args, **kwargs):
 
 
 @contextlib.contextmanager
-def iarange(name, size, subsample_size=0):
+def iarange(name, size, subsample_size=0, subsample=None):
     """
     Context manager for ranges indexing iid variables, optionally subsampling.
 
@@ -131,6 +131,9 @@ def iarange(name, size, subsample_size=0):
     :param str name: A name that will be used for this site in a Trace.
     :param int size: The size of the collection being subsampled (like `stop` in builtin `range`).
     :param int subsample_size: Size of minibatches used in subsampling. Defaults to `size` if set to 0.
+    :param subsample: Optional custom subsample for user-defined subsampling schemes.
+        If specified, then `subsample_size` will be set to `len(subsample)`.
+    :type subsample: Anything supporting `len()`.
     :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
 
     Examples::
@@ -138,20 +141,29 @@ def iarange(name, size, subsample_size=0):
         # This version is vectorized:
         >>> with iarange('data', 100, subsample_size=10) as batch:
                 observe('obs', normal, data.index_select(0, batch), mu, sigma)
+
         # This version manually iterates through the batch to deal with control flow.
         >>> with iarange('data', 100, subsample_size=10) as batch:
                 for i in batch:
                     if z[i]:  # Prevents vectorization.
                         observe('obs_{}'.format(i), normal, data[i], mu, sigma)
+
+        # This wraps a user-defined subsampling method for use in pyro:
+        >>> with iarange('data', 100, subsample=my_custom_subsample) as batch:
+                assert batch is my_custom_subsample
+                observe('obs', normal, data.index_select(0, batch), mu, sigma)
     """
-    if subsample_size == 0 or subsample_size >= size:
+    if subsample is not None:
+        subsample_size = len(subsample)
+        assert subsample_size <= size, 'subsample is larger than size'
+    elif subsample_size == 0 or subsample_size >= size:
         subsample_size = size
     if subsample_size == size:
         # If not subsampling, there is no need to scale and we can ignore the _PYRO_STACK.
         yield Variable(torch.LongTensor(list(range(size))))
         return
 
-    subsample = sample(name, Subsample(size, subsample_size))
+    subsample = sample(name, Subsample(size, subsample_size), obs=subsample)
     if len(_PYRO_STACK) == 0:
         yield subsample
     else:
@@ -161,9 +173,17 @@ def iarange(name, size, subsample_size=0):
             yield subsample
 
 
-def irange(name, size, subsample_size=0):
+def irange(name, size, subsample_size=0, subsample=None):
     """
-    Non-vectorized version of iarange.
+    Non-vectorized version of `iarange`. See `iarange` for details.
+
+    :param str name: A name that will be used for this site in a Trace.
+    :param int size: The size of the collection being subsampled (like `stop` in builtin `range`).
+    :param int subsample_size: Size of minibatches used in subsampling. Defaults to `size` if set to 0.
+    :param subsample: Optional custom subsample for user-defined subsampling schemes.
+        If specified, then `subsample_size` will be set to `len(subsample)`.
+    :type subsample: Anything supporting `len()`.
+    :return: A context manager yielding a single 1-dimensional `torch.Tensor` of indices.
 
     Examples::
 
@@ -171,7 +191,10 @@ def irange(name, size, subsample_size=0):
                 if z[i]:  # Prevents vectorization.
                     observe('obs_{}'.format(i), normal, data[i], mu, sigma)
     """
-    with iarange(name, size, subsample_size) as batch:
+    if subsample is not None:
+        subsample_size = len(subsample)
+
+    with iarange(name, size, subsample_size, subsample) as batch:
         # Wrap computation in an independence context.
         indep_context = LambdaPoutine(None, name, 1.0, 'list', 0, subsample_size)
         for i in batch.data:
diff --git a/tests/poutine/test_mapdata.py b/tests/poutine/test_mapdata.py
index 3c01014c27..83c2465c35 100644
--- a/tests/poutine/test_mapdata.py
+++ b/tests/poutine/test_mapdata.py
@@ -160,27 +160,65 @@ def model(means, stds):
             assert tr.nodes[name]["scale"] == 4.0 * 2.0
 
 
-def test_replay_iarange():
-    pyro.set_rng_seed(0)
+def iarange_model():
+    with pyro.iarange('iarange', 20, 5) as batch:
+        result = list(batch.data)
+    return result
 
-    def model():
-        with pyro.iarange('iarange', 10, 2) as batch:
-            return list(batch.data)
 
-    traced_model = poutine.trace(model)
-    expected = traced_model()
-    actual = poutine.replay(model, traced_model.trace)()
-    assert actual == expected
+def irange_model():
+    result = []
+    for i in pyro.irange('irange', 20, 5):
+        result.append(i)
+    return result
+
+
+def map_data_model():
+    return pyro.map_data('mapdata', range(20), lambda i, x: i, batch_size=5)
 
 
-def test_replay_map_data():
+@pytest.mark.parametrize('model', [
+    iarange_model,
+    irange_model,
+    map_data_model,
+], ids=['iarange', 'irange', 'map_data'])
+def test_replay(model):
     pyro.set_rng_seed(0)
-    data = list(range(10))
 
-    def model():
-        return pyro.map_data('mapdata', data, lambda i, x: i, batch_size=2)
+    traced_model = poutine.trace(model)
+    original = traced_model()
+
+    replayed = poutine.replay(model, traced_model.trace)()
+    assert replayed == original
+
+    different = traced_model()
+    assert different != original
+
+
+def iarange_custom_model(subsample):
+    with pyro.iarange('iarange', 20, subsample=subsample) as batch:
+        result = list(batch.data)
+    return result
+
+
+def irange_custom_model(subsample):
+    result = []
+    for i in pyro.irange('irange', 20, subsample=subsample):
+        result.append(i)
+    return result
+
+
+@pytest.mark.xfail
+@pytest.mark.parametrize('model', [iarange_custom_model, irange_custom_model],
+                         ids=['iarange', 'irange'])
+def test_replay_custom(model):
+    pyro.set_rng_seed(0)
 
     traced_model = poutine.trace(model)
-    expected = traced_model()
-    actual = poutine.replay(model, traced_model.trace)()
-    assert actual == expected
+    original = traced_model([1, 2, 3, 4])
+
+    replayed = poutine.replay(model, traced_model.trace)([5, 6, 7, 8])
+    assert replayed == original
+
+    different = traced_model([5, 6, 7, 8])
+    assert different != original


Code review comments:
diff_hunk: @@ -160,27 +160,65 @@ def model(means, stds):
             assert tr.nodes[name]["scale"] == 4.0 * 2.0
 
 
-def test_replay_iarange():
-    pyro.set_rng_seed(0)
+def iarange_model():
+    with pyro.iarange('iarange', 20, 5) as batch:
+        result = list(batch.data)
+    return result
 
-    def model():
-        with pyro.iarange('iarange', 10, 2) as batch:
-            return list(batch.data)
 
-    traced_model = poutine.trace(model)
-    expected = traced_model()
-    actual = poutine.replay(model, traced_model.trace)()
-    assert actual == expected
+def irange_model():
+    result = []
+    for i in pyro.irange('irange', 20, 5):
+        result.append(i)
+    return result
+
+
+def map_data_model():
+    return pyro.map_data('mapdata', range(20), lambda i, x: i, batch_size=5)
 
 
-def test_replay_map_data():
+@pytest.mark.parametrize('model', [
+    iarange_model,
+    irange_model,
+    map_data_model,
+], ids=['iarange', 'irange', 'map_data'])
+def test_replay(model):
     pyro.set_rng_seed(0)
-    data = list(range(10))
 
-    def model():
-        return pyro.map_data('mapdata', data, lambda i, x: i, batch_size=2)
+    traced_model = poutine.trace(model)
+    original = traced_model()
+
+    replayed = poutine.replay(model, traced_model.trace)()
+    assert replayed == original
+
+    different = traced_model()
+    assert different != original
+
+
+def iarange_custom_model(subsample):
+    with pyro.iarange('iarange', 20, subsample=subsample) as batch:
+        result = list(batch.data)
+    return result
+
+
+def irange_custom_model(subsample):
+    result = []
+    for i in pyro.irange('irange', 20, subsample=subsample):
+        result.append(i)
+    return result
+
+
+@pytest.mark.xfail
path: tests/poutine/test_mapdata.py
text: @author:
Tests currently fail.