Clear .unconstrained weakrefs before pickling; rebuild them more often

pyro/contrib/easyguide/easyguide.py

@@ -230,7 +230,7 @@ def __init__(self, guide, sites):
         self.event_shape = torch.Size([sum(self._site_sizes.values())])
 
     def __getstate__(self):
-        state = self.__dict__.copy()
+        state = getattr(super(), "__getstate__", self.__dict__.copy)()
         state["_guide"] = state["_guide"]()  # weakref -> ref
         return state
 

pyro/contrib/util.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2017-2019 Uber Technologies, Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+import weakref
 from collections import OrderedDict
 
 import torch
@@ -75,3 +76,22 @@ def iter_plates_to_shape(shape):
     # Go backwards (right to left)
     for i, s in enumerate(shape[::-1]):
         yield pyro.plate("plate_" + str(i), s)
+
+
+def check_no_weakref(obj, path="", avoid_ids=None):
+    """Attempts to check that an object has no weakrefs."""
+    if avoid_ids is None:
+        avoid_ids = {id(obj)}
+
+    if isinstance(obj, weakref.ref):
+        raise ValueError(f"Weakref found at {path}")
+    elif isinstance(obj, dict):
+        for k, v in obj.items():
+            if id(v) not in avoid_ids:
+                check_no_weakref(v, path + f"[{k}]")
+    elif isinstance(obj, (list, tuple)):
+        for i, v in enumerate(obj):
+            if id(v) not in avoid_ids:
+                check_no_weakref(v, path + f"[{i}]")
+    elif hasattr(obj, "__dict__"):
+        check_no_weakref(obj.__dict__, path)

pyro/distributions/torch_patch.py

@@ -44,13 +44,12 @@ def decorator(new_fn):
 # TODO: Move upstream to allow for pickle serialization of transforms
 @patch_dependency("torch.distributions.transforms.Transform.__getstate__")
 def _Transform__getstate__(self):
-    attrs = {}
-    for k, v in self.__dict__.items():
+    super_ = super(torch.distributions.transforms.Transform, self)
+    state = getattr(super_, "__getstate__", self.__dict__.copy)()
+    for k, v in state.items():
         if isinstance(v, weakref.ref):
-            attrs[k] = None
-        else:
-            attrs[k] = v
-    return attrs
+            state[k] = None
+    return state
 
 
 # TODO move upstream

pyro/infer/elbo.py

@@ -66,9 +66,10 @@ class ELBO(object, metaclass=ABCMeta):
         behave unexpectedly relative to standard PyTorch when working with
         :class:`~pyro.nn.PyroModule` s.
 
-        Users are therefore strongly encouraged to use this interface in conjunction
-        with :func:`~pyro.enable_module_local_param` which will override the default
-        implicit sharing of parameters across :class:`~pyro.nn.PyroModule` instances.
+        Users are therefore strongly encouraged to use this interface in
+        conjunction with ``pyro.settings.set(module_local_params=True)`` which
+        will override the default implicit sharing of parameters across
+        :class:`~pyro.nn.PyroModule` instances.
 
     :param num_particles: The number of particles/samples used to form the ELBO
         (gradient) estimators.

pyro/nn/module.py

@@ -14,6 +14,7 @@
 """
 import functools
 import inspect
+import weakref
 from collections import OrderedDict, namedtuple
 
 import torch
@@ -503,7 +504,9 @@ def __getattr__(self, name):
                         _PYRO_PARAM_STORE._param_to_name[unconstrained_value] = fullname
                     return pyro.param(fullname, event_dim=event_dim)
                 else:  # Cannot determine supermodule and hence cannot compute fullname.
-                    return transform_to(constraint)(unconstrained_value)
+                    constrained_value = transform_to(constraint)(unconstrained_value)
+                    constrained_value.unconstrained = weakref.ref(unconstrained_value)
+                    return constrained_value
 
         # PyroSample trigger pyro.sample statements.
         if "_pyro_samples" in self.__dict__:
@@ -661,6 +664,12 @@ def __delattr__(self, name):
 
         super().__delattr__(name)
 
+    def __getstate__(self):
+        # Remove weakrefs in preparation for pickling.
+        for param in self.parameters(recurse=True):
+            param.__dict__.pop("unconstrained", None)
+        return getattr(super(), "__getstate__", self.__dict__.copy)()
+
 
 def pyro_method(fn):
     """

pyro/params/param_store.py

@@ -243,10 +243,11 @@ def get_state(self) -> dict:
         """
         Get the ParamStore state.
         """
-        state = {
-            "params": self._params.copy(),
-            "constraints": self._constraints.copy(),
-        }
+        params = self._params.copy()
+        # Remove weakrefs in preparation for pickling.
+        for param in params.values():
+            param.__dict__.pop("unconstrained", None)
+        state = {"params": params, "constraints": self._constraints.copy()}
         return state
 
     def set_state(self, state: dict):

pyro/poutine/guide.py

@@ -30,6 +30,12 @@ def __init__(self, model: Callable):
     def model(self):
         return self._model[0]
 
+    def __getstate__(self):
+        # Avoid pickling the trace.
+        state = super().__getstate__()
+        state.pop("trace")
+        return state
+
     def __call__(self, *args, **kwargs) -> Dict[str, torch.Tensor]:
         """
         Draws posterior samples from the guide and replays the model against

tests/nn/test_module.py

@@ -374,7 +374,9 @@ def forward(self):
 
         # mutate
         for _, x in m.named_pyro_params():
-            x.unconstrained().data += torch.randn(())
+            if hasattr(x, "unconstrained"):
+                x = x.unconstrained()
+            x.data += torch.randn(x.shape)
         state1 = m()
         for x, y in zip(state0, state1):
             assert not (x == y).all()