Repository: pyro-ppl/pyro

Issue:
Graph retention in JitTraceGraph_ELBO
In reviewing #1891, @martinjankowiak  noted that `JitTraceGraph_ELBO` isn't using `retain_graph`. It could make sense to simply switch to `backward(retain_graph=self.retain_graph)`. I'm reluctant to do that without some understanding of how graph retention interacts with jit.


Discussion:

Resolving patch:
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
@@ -91,7 +91,7 @@ def _compute_downstream_costs(model_trace, guide_trace,  #
     return downstream_costs, downstream_guide_cost_nodes
 
 
-def _compute_elbo_reparam(model_trace, guide_trace, non_reparam_nodes):
+def _compute_elbo_reparam(model_trace, guide_trace):
     elbo = 0.0
     surrogate_elbo = 0.0
 
@@ -219,37 +219,49 @@ def loss_and_grads(self, model, guide, *args, **kwargs):
         Performs backward on the latter. Num_particle many samples are used to form the estimators.
         If baselines are present, a baseline loss is also constructed and differentiated.
         """
+
+        loss, surrogate_loss = self._loss_and_surrogate_loss(model, guide, *args, **kwargs)
+
+        torch_backward(surrogate_loss, retain_graph=self.retain_graph)
+
+        loss = torch_item(loss)
+        warn_if_nan(loss, "loss")
+        return loss
+
+    def _loss_and_surrogate_loss(self, model, guide, *args, **kwargs):
+
         loss = 0.0
-        weight = 1./self.num_particles
+        surrogate_loss = 0.0
+
         for model_trace, guide_trace in self._get_traces(model, guide, *args, **kwargs):
-            loss += self._loss_and_grads_particle(weight, model_trace, guide_trace)
-        return loss
 
-    def _loss_and_grads_particle(self, weight, model_trace, guide_trace):
+            lp, slp = self._loss_and_surrogate_loss_particle(model_trace, guide_trace, *args, **kwargs)
+            loss += lp
+            surrogate_loss += slp
+
+        loss /= self.num_particles
+        surrogate_loss /= self.num_particles
+
+        return loss, surrogate_loss
+
+    def _loss_and_surrogate_loss_particle(self, model_trace, guide_trace, *args, **kwargs):
+
         # compute elbo for reparameterized nodes
-        non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
-        elbo, surrogate_elbo = _compute_elbo_reparam(model_trace, guide_trace, non_reparam_nodes)
+        elbo, surrogate_elbo = _compute_elbo_reparam(model_trace, guide_trace)
+        baseline_loss = 0.0
 
         # the following computations are only necessary if we have non-reparameterizable nodes
-        baseline_loss = 0.0
+        non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
         if non_reparam_nodes:
             downstream_costs, _ = _compute_downstream_costs(model_trace, guide_trace, non_reparam_nodes)
             surrogate_elbo_term, baseline_loss = _compute_elbo_non_reparam(guide_trace,
-                                                                           non_reparam_nodes, downstream_costs)
+                                                                           non_reparam_nodes,
+                                                                           downstream_costs)
             surrogate_elbo += surrogate_elbo_term
 
-        # collect parameters to train from model and guide
-        trainable_params = any(site["type"] == "param"
-                               for trace in (model_trace, guide_trace)
-                               for site in trace.nodes.values())
-
-        if trainable_params:
-            surrogate_loss = -surrogate_elbo
-            torch_backward(weight * (surrogate_loss + baseline_loss), retain_graph=self.retain_graph)
+        surrogate_loss = -surrogate_elbo + baseline_loss
 
-        loss = -torch_item(elbo)
-        warn_if_nan(loss, "loss")
-        return weight * loss
+        return elbo, surrogate_loss
 
 
 class JitTraceGraph_ELBO(TraceGraph_ELBO):
@@ -270,43 +282,24 @@ class JitTraceGraph_ELBO(TraceGraph_ELBO):
     def loss_and_grads(self, model, guide, *args, **kwargs):
         kwargs['_pyro_model_id'] = id(model)
         kwargs['_pyro_guide_id'] = id(guide)
-        if getattr(self, '_loss_and_surrogate_loss', None) is None:
+        if getattr(self, '_jit_loss_and_surrogate_loss', None) is None:
             # build a closure for loss_and_surrogate_loss
             weakself = weakref.ref(self)
 
             @pyro.ops.jit.trace(ignore_warnings=self.ignore_jit_warnings,
                                 jit_options=self.jit_options)
-            def loss_and_surrogate_loss(*args, **kwargs):
+            def jit_loss_and_surrogate_loss(*args, **kwargs):
                 kwargs.pop('_pyro_model_id')
                 kwargs.pop('_pyro_guide_id')
                 self = weakself()
-                loss = 0.0
-                surrogate_loss = 0.0
-                weight = 1.0 / self.num_particles
-                for model_trace, guide_trace in self._get_traces(model, guide, *args, **kwargs):
-                    # compute elbo for reparameterized nodes
-                    non_reparam_nodes = set(guide_trace.nonreparam_stochastic_nodes)
-                    elbo, surrogate_elbo = _compute_elbo_reparam(model_trace, guide_trace, non_reparam_nodes)
-
-                    # the following computations are only necessary if we have non-reparameterizable nodes
-                    baseline_loss = 0.0
-                    if non_reparam_nodes:
-                        downstream_costs, _ = _compute_downstream_costs(model_trace, guide_trace, non_reparam_nodes)
-                        surrogate_elbo_term, baseline_loss = _compute_elbo_non_reparam(guide_trace,
-                                                                                       non_reparam_nodes,
-                                                                                       downstream_costs)
-                        surrogate_elbo += surrogate_elbo_term
-
-                    loss = loss - weight * elbo
-                    surrogate_loss = surrogate_loss - weight * surrogate_elbo
-
-                return loss, surrogate_loss
-
-            self._loss_and_surrogate_loss = loss_and_surrogate_loss
-
-        loss, surrogate_loss = self._loss_and_surrogate_loss(*args, **kwargs)
-        surrogate_loss.backward()
-        loss = loss.item()
+                return self._loss_and_surrogate_loss(model, guide, *args, **kwargs)
+
+            self._jit_loss_and_surrogate_loss = jit_loss_and_surrogate_loss
 
+        loss, surrogate_loss = self._jit_loss_and_surrogate_loss(*args, **kwargs)
+
+        surrogate_loss.backward(retain_graph=self.retain_graph)  # triggers jit compilation
+
+        loss = loss.item()
         warn_if_nan(loss, "loss")
         return loss
diff --git a/pyro/infer/util.py b/pyro/infer/util.py
--- a/pyro/infer/util.py
+++ b/pyro/infer/util.py
@@ -37,9 +37,9 @@ def torch_item(x):
 def torch_backward(x, retain_graph=None):
     """
     Like ``x.backward()`` for a :class:`~torch.Tensor`, but also accepts
-    numbers (a no-op if given a number).
+    numbers and tensors without grad_fn (resulting in a no-op)
     """
-    if torch.is_tensor(x):
+    if torch.is_tensor(x) and x.grad_fn:
         x.backward(retain_graph=retain_graph)
 
 


Test patch: