Repository: pyro-ppl/pyro

Issue:
Memory leak due to cyclic references in Trace
@karalets noticed memory growth when using Pyro.

We tracked the number of non-garbage-collected objects during SVI inference. We found that objects grow over iterations but are occasionally collected. This evidence points to a cyclic reference somewhere. Among the amassing objects were `Variable`s and  Pyro `Trace` objects, suggesting that the cyclic reference is either in Pyro internals or in the networkx `DiGraph` object from which `Trace` inherits.


Discussion:

Resolving patch:
diff --git a/pyro/poutine/trace.py b/pyro/poutine/trace.py
--- a/pyro/poutine/trace.py
+++ b/pyro/poutine/trace.py
@@ -20,13 +20,15 @@ def _warn_if_nan(name, value):
     # Note that -inf log_pdf is fine: it is merely a zero-probability event.
 
 
-class Trace(networkx.DiGraph):
+class DiGraph(networkx.DiGraph):
+    node_dict_factory = collections.OrderedDict
+
+
+class Trace(object):
     """
     Execution trace data structure
     """
 
-    node_dict_factory = collections.OrderedDict
-
     def __init__(self, *args, **kwargs):
         """
         :param string graph_type: string specifying the kind of trace graph to construct
@@ -34,12 +36,59 @@ def __init__(self, *args, **kwargs):
         Constructor. Currently identical to networkx.``DiGraph(\*args, \**kwargs)``,
         except for storing the graph_type attribute
         """
+        self._graph = DiGraph(*args, **kwargs)
         graph_type = kwargs.pop("graph_type", "flat")
         assert graph_type in ("flat", "dense"), \
             "{} not a valid graph type".format(graph_type)
         self.graph_type = graph_type
         super(Trace, self).__init__(*args, **kwargs)
 
+    def __del__(self):
+        # Work around cyclic reference bugs in networkx.DiGraph
+        # See https://github.com/uber/pyro/issues/798
+        self._graph.__dict__.clear()
+
+    @property
+    def nodes(self):
+        return self._graph.nodes
+
+    @property
+    def edges(self):
+        return self._graph.edges
+
+    @property
+    def graph(self):
+        return self._graph.graph
+
+    @property
+    def remove_node(self):
+        return self._graph.remove_node
+
+    @property
+    def add_edge(self):
+        return self._graph.add_edge
+
+    @property
+    def is_directed(self):
+        return self._graph.is_directed
+
+    @property
+    def in_degree(self):
+        return self._graph.in_degree
+
+    @property
+    def successors(self):
+        return self._graph.successors
+
+    def __contains__(self, site_name):
+        return site_name in self._graph
+
+    def __iter__(self):
+        return iter(self._graph)
+
+    def __len__(self):
+        return len(self._graph)
+
     def add_node(self, site_name, *args, **kwargs):
         """
         :param string site_name: the name of the site to be added
@@ -56,7 +105,7 @@ def add_node(self, site_name, *args, **kwargs):
                 "site {} already in trace".format(site_name)
 
         # XXX should copy in case site gets mutated, or dont bother?
-        super(Trace, self).add_node(site_name, *args, **kwargs.copy())
+        self._graph.add_node(site_name, *args, **kwargs)
 
     def copy(self):
         """
@@ -64,9 +113,9 @@ def copy(self):
         Identical to super(Trace, self).copy(), but preserves the type
         and the self.graph_type attribute
         """
-        trace = super(Trace, self).copy()
+        trace = Trace()
+        trace._graph = self._graph.copy()
         trace.graph_type = self.graph_type
-        trace.__class__ = Trace
         return trace
 
     def log_pdf(self, site_filter=lambda name, site: True):


Test patch:
diff --git a/tests/test_memory.py b/tests/test_memory.py
new file mode 100644
--- /dev/null
+++ b/tests/test_memory.py
@@ -0,0 +1,156 @@
+from __future__ import absolute_import, division, print_function
+
+import gc
+
+import networkx as nx
+import torch
+from torch.autograd import Variable
+
+import pyro
+import pyro.distributions as dist
+from pyro import poutine
+from pyro.infer.svi import SVI
+from pyro.optim import Adam
+from pyro.poutine.trace import Trace
+
+
+def count_objects_of_type(type_):
+    return sum(1 for obj in gc.get_objects() if isinstance(obj, type_))
+
+
+def test_trace():
+    n = 11
+    data = Variable(torch.zeros(n))
+
+    def model(data):
+        loc = pyro.param('loc', Variable(torch.zeros(n), requires_grad=True))
+        scale = pyro.param('log_scale', Variable(torch.zeros(n), requires_grad=True)).exp()
+        pyro.sample('obs', dist.Normal(loc, scale), obs=data)
+
+    counts = []
+    gc.collect()
+    gc.collect()
+    expected = count_objects_of_type(Trace)
+    for _ in range(10):
+        poutine.trace(model)(data)
+        counts.append(count_objects_of_type(Trace))
+
+    assert set(counts) == set([expected]), counts
+
+
+class Foo(dict):
+    pass
+
+
+def test_dict_copy():
+    counts = []
+    gc.collect()
+    gc.collect()
+    f = Foo()
+    for _ in range(10):
+        f.copy()
+        counts.append(count_objects_of_type(Foo))
+
+    assert set(counts) == set([1]), counts
+
+
+def test_networkx_copy():
+    counts = []
+    gc.collect()
+    gc.collect()
+    g = nx.DiGraph()
+    expected = count_objects_of_type(nx.DiGraph)
+    for _ in range(10):
+        h = g.fresh_copy()
+        h.__dict__.clear()
+        del h
+        counts.append(count_objects_of_type(nx.DiGraph))
+
+    assert set(counts) == set([expected]), counts
+
+
+def test_copy():
+    counts = []
+    gc.collect()
+    gc.collect()
+    tr = Trace()
+    expected = count_objects_of_type(Trace)
+    for _ in range(10):
+        tr.copy()
+        counts.append(count_objects_of_type(Trace))
+
+    assert set(counts) == set([expected]), counts
+
+
+def test_trace_copy():
+    n = 11
+    data = Variable(torch.zeros(n))
+
+    def model(data):
+        loc = pyro.param('loc', Variable(torch.zeros(n), requires_grad=True))
+        scale = pyro.param('log_scale', Variable(torch.zeros(n), requires_grad=True)).exp()
+        pyro.sample('obs', dist.Normal(loc, scale), obs=data)
+
+    counts = []
+    gc.collect()
+    gc.collect()
+    expected = count_objects_of_type(Trace)
+    for _ in range(10):
+        poutine.trace(model).get_trace(data).copy()
+        counts.append(count_objects_of_type(Trace))
+
+    assert set(counts) == set([expected]), counts
+
+
+def trace_replay(model, guide, *args):
+    guide_trace = poutine.trace(guide).get_trace(*args)
+    poutine.trace(poutine.replay(model, guide_trace)).get_trace(*args)
+
+
+def test_trace_replay():
+    n = 11
+    data = Variable(torch.zeros(n))
+
+    def model(data):
+        loc = pyro.param('loc', Variable(torch.zeros(n), requires_grad=True))
+        scale = pyro.param('log_scale', Variable(torch.zeros(n), requires_grad=True)).exp()
+        pyro.sample('obs', dist.Normal(loc, scale), obs=data)
+
+    def guide(data):
+        pass
+
+    counts = []
+    gc.collect()
+    gc.collect()
+    expected = count_objects_of_type(Trace)
+    for _ in range(10):
+        trace_replay(model, guide, data)
+        counts.append(count_objects_of_type(Trace))
+
+    assert set(counts) == set([expected]), counts
+
+
+def test_svi():
+    n = 11
+    data = Variable(torch.zeros(n))
+
+    def model(data):
+        loc = pyro.param('loc', Variable(torch.zeros(n), requires_grad=True))
+        scale = pyro.param('log_scale', Variable(torch.zeros(n), requires_grad=True)).exp()
+        pyro.sample('obs', dist.Normal(loc, scale), obs=data)
+
+    def guide(data):
+        pass
+
+    optim = Adam({'lr': 1e-3})
+    inference = SVI(model, guide, optim, 'ELBO')
+
+    counts = []
+    gc.collect()
+    gc.collect()
+    expected = count_objects_of_type(Trace)
+    for _ in range(10):
+        inference.step(data)
+        counts.append(count_objects_of_type(Trace))
+
+    assert set(counts) == set([expected]), counts