Repository: pyro-ppl/pyro

Issue:
enum_discrete segfaults on Pytorch 0.2.0 release (ok on master)



Discussion:
This appears to be caused by a large number `surrogate_elbo_particle` terms being added together, and has caused a segfault even when `enum_discrete=False` and `num_particles=20` in `TraceGraph_ELBO` (in the test `tests/infer/test_enum.py::test_bern_elbo_gradient`).
@martinjankowiak observe that in final Beta-Bernoulli example of the [SVI Part I](https://github.com/uber/pyro/blob/dev/tutorial/source/svi_part_i.ipynb) tutorial, pyro currenly works with `num_particles=1` but segfaults with `num_particles=100`, yet there was no segfault a couple weeks ago. This might be a good example for git-bisecting. 
```diff
- svi = SVI(model, guide, optimizer, loss="ELBO", num_particles=1)
+ svi = SVI(model, guide, optimizer, loss="ELBO", num_particles=100)
```

Resolving patch:
diff --git a/examples/ss_vae_M2.py b/examples/ss_vae_M2.py
--- a/examples/ss_vae_M2.py
+++ b/examples/ss_vae_M2.py
@@ -14,9 +14,9 @@
 import torch.nn as nn
 
 version_warning = '''
-11/02/2017: This example does not work with the release version 0.2 of pytorch
-Please install Pytorch from the latest master branch of pytorch or wait a week for the new release
-This example may fail due to https://github.com/uber/pyro/issues/377
+11/02/2017: This example does not work with the release version 0.2 of pytorch.
+Please install Pytorch from the latest master branch of pytorch or wait a week for the new release.
+This example uses a data loader that requires very recent PyTorch features.
 '''
 torch_version = pyro.util.parse_torch_version()
 if (torch_version < (0, 2, 1) and not torch_version[-1].startswith("+")):
diff --git a/pyro/infer/trace_elbo.py b/pyro/infer/trace_elbo.py
--- a/pyro/infer/trace_elbo.py
+++ b/pyro/infer/trace_elbo.py
@@ -129,8 +129,6 @@ def loss_and_grads(self, model, guide, *args, **kwargs):
         Performs backward on the latter. Num_particle many samples are used to form the estimators.
         """
         elbo = 0.0
-        surrogate_elbo = 0.0
-        trainable_params = set()
         # grab a trace from the generator
         for weight, model_trace, guide_trace, log_r in self._get_traces(model, guide, *args, **kwargs):
             elbo_particle = weight * 0
@@ -163,23 +161,19 @@ def loss_and_grads(self, model, guide, *args, **kwargs):
                 surrogate_elbo_particle[weight_eq_zero] = 0.0
 
             elbo += torch_data_sum(weight * elbo_particle)
-            surrogate_elbo += torch_sum(weight * surrogate_elbo_particle)
+            surrogate_elbo_particle = torch_sum(weight * surrogate_elbo_particle)
 
-            # grab model parameters to train
-            for name in model_trace.nodes.keys():
-                if model_trace.nodes[name]["type"] == "param":
-                    trainable_params.add(model_trace.nodes[name]["value"])
+            # collect parameters to train from model and guide
+            trainable_params = set(site["value"]
+                                   for trace in (model_trace, guide_trace)
+                                   for site in trace.nodes.values()
+                                   if site["type"] == "param")
 
-            # grab guide parameters to train
-            for name in guide_trace.nodes.keys():
-                if guide_trace.nodes[name]["type"] == "param":
-                    trainable_params.add(guide_trace.nodes[name]["value"])
+            if trainable_params:
+                surrogate_loss_particle = -surrogate_elbo_particle
+                torch_backward(surrogate_loss_particle)
+                pyro.get_param_store().mark_params_active(trainable_params)
 
         loss = -elbo
-        surrogate_loss = -surrogate_elbo
-        if trainable_params:
-            torch_backward(surrogate_loss)
-
-        pyro.get_param_store().mark_params_active(trainable_params)
 
         return loss
diff --git a/pyro/infer/tracegraph_elbo.py b/pyro/infer/tracegraph_elbo.py
--- a/pyro/infer/tracegraph_elbo.py
+++ b/pyro/infer/tracegraph_elbo.py
@@ -98,7 +98,6 @@ def loss_and_grads(self, model, guide, *args, **kwargs):
         If baselines are present, a baseline loss is also constructed and differentiated.
         """
         elbo = 0.0
-        trainable_params = set()
 
         for weight, model_trace, guide_trace in self._get_traces(model, guide, *args, **kwargs):
 
@@ -280,23 +279,19 @@ def get_baseline_options(site_baseline):
                 surrogate_elbo_particle += weight * elbo_reinforce_terms_particle
                 torch_backward(baseline_loss_particle)
 
-            # grab model parameters to train
-            for name in model_trace.nodes.keys():
-                if model_trace.nodes[name]["type"] == "param":
-                    trainable_params.add(model_trace.nodes[name]["value"])
-
-            # grab guide parameters to train
-            for name in guide_trace.nodes.keys():
-                if guide_trace.nodes[name]["type"] == "param":
-                    trainable_params.add(guide_trace.nodes[name]["value"])
-
-            # mark all params seen in trace as active so that gradient steps are taken downstream
-            pyro.get_param_store().mark_params_active(trainable_params)
+            # collect parameters to train from model and guide
+            trainable_params = set(site["value"]
+                                   for trace in (model_trace, guide_trace)
+                                   for site in trace.nodes.values()
+                                   if site["type"] == "param")
 
             surrogate_loss_particle = -surrogate_elbo_particle
             if trainable_params:
                 torch_backward(surrogate_loss_particle)
 
+                # mark all params seen in trace as active so that gradient steps are taken downstream
+                pyro.get_param_store().mark_params_active(trainable_params)
+
         loss = -elbo
 
         return loss


Test patch:
diff --git a/tests/common.py b/tests/common.py
--- a/tests/common.py
+++ b/tests/common.py
@@ -16,8 +16,6 @@
 from pytest import approx
 from torch.autograd import Variable
 
-import pyro.util
-
 torch.set_default_tensor_type(os.environ.get('PYRO_TENSOR_TYPE', 'torch.DoubleTensor'))
 
 """
@@ -230,14 +228,6 @@ def assert_not_equal(x, y, prec=1e-5, msg=''):
     raise AssertionError("{} \nValues are equal: x={}, y={}, prec={}".format(msg, x, y, prec))
 
 
-# Avoid tests that fail due to https://github.com/uber/pyro/issues/377
-torch_version = pyro.util.parse_torch_version()
-segfaults_on_pytorch_020 = pytest.mark.skipif(
-    torch_version < (0, 2, 1) and not torch_version[-1].startswith("+"),
-    reason=("pytorch segfaults at 0.2.0_4, fixed by 0.2.0+f964105. "
-            "See https://github.com/uber/pyro/issues/377"))
-
-
 class TestCase(unittest.TestCase):
     precision = 1e-5
 
diff --git a/tests/infer/test_enum.py b/tests/infer/test_enum.py
--- a/tests/infer/test_enum.py
+++ b/tests/infer/test_enum.py
@@ -14,7 +14,7 @@
 from pyro.infer.enum import iter_discrete_traces
 from pyro.infer.trace_elbo import Trace_ELBO
 from pyro.infer.tracegraph_elbo import TraceGraph_ELBO
-from tests.common import assert_equal, segfaults_on_pytorch_020, xfail_if_not_implemented
+from tests.common import assert_equal, xfail_if_not_implemented
 
 
 @pytest.mark.parametrize("graph_type", ["flat", "dense"])
@@ -162,7 +162,6 @@ def test_gmm_batch_iter_discrete_traces(model, data_size, graph_type):
     assert len(traces) == 2
 
 
-@segfaults_on_pytorch_020
 @pytest.mark.parametrize("trace_graph", [False, True], ids=["dense", "flat"])
 @pytest.mark.parametrize("model,guide", [
     (gmm_model, gmm_guide),
@@ -200,7 +199,6 @@ def finite_difference(eval_loss, delta=0.1):
     return grads
 
 
-@segfaults_on_pytorch_020
 @pytest.mark.parametrize("enum_discrete", [True, False], ids=["sum", "sample"])
 @pytest.mark.parametrize("trace_graph", [False, True], ids=["dense", "flat"])
 def test_bern_elbo_gradient(enum_discrete, trace_graph):
@@ -235,7 +233,6 @@ def guide():
     assert_equal(actual_grads, expected_grads, prec=0.1)
 
 
-@segfaults_on_pytorch_020
 @pytest.mark.parametrize("model,guide", [
     (gmm_model, gmm_guide),
     (gmm_batch_model, gmm_batch_guide),
diff --git a/tests/infer/test_valid_models.py b/tests/infer/test_valid_models.py
--- a/tests/infer/test_valid_models.py
+++ b/tests/infer/test_valid_models.py
@@ -10,7 +10,6 @@
 import pyro.distributions as dist
 from pyro.infer import SVI
 from pyro.optim import Adam
-from tests.common import segfaults_on_pytorch_020
 
 # This file tests a variety of model,guide pairs with valid and invalid structure.
 
@@ -396,7 +395,6 @@ def guide():
     assert_error(model, guide, trace_graph=True)
 
 
-@segfaults_on_pytorch_020
 def test_enum_discrete_single_ok():
 
     def model():
@@ -410,7 +408,6 @@ def guide():
     assert_ok(model, guide, enum_discrete=True)
 
 
-@segfaults_on_pytorch_020
 def test_enum_discrete_single_single_ok():
 
     def model():
@@ -426,7 +423,6 @@ def guide():
     assert_ok(model, guide, enum_discrete=True)
 
 
-@segfaults_on_pytorch_020
 def test_enum_discrete_irange_single_ok():
 
     def model():
@@ -442,7 +438,6 @@ def guide():
     assert_ok(model, guide, enum_discrete=True)
 
 
-@segfaults_on_pytorch_020
 def test_iarange_enum_discrete_batch_ok():
 
     def model():
@@ -458,7 +453,6 @@ def guide():
     assert_ok(model, guide, enum_discrete=True)
 
 
-@segfaults_on_pytorch_020
 def test_iarange_enum_discrete_no_discrete_vars_ok():
 
     def model():
@@ -476,7 +470,6 @@ def guide():
     assert_ok(model, guide, enum_discrete=True)
 
 
-@segfaults_on_pytorch_020
 def test_no_iarange_enum_discrete_batch_error():
 
     def model():
@@ -490,7 +483,6 @@ def guide():
     assert_error(model, guide, enum_discrete=True)
 
 
-@segfaults_on_pytorch_020
 def test_enum_discrete_global_local_error():
 
     def model():