Add an option to stop PyroModules from sharing parameters

pyro/__init__.py

@@ -9,12 +9,14 @@
     barrier,
     clear_param_store,
     deterministic,
+    enable_module_local_param,
     enable_validation,
     factor,
     get_param_store,
     iarange,
     irange,
     module,
+    module_local_param_enabled,
     param,
     plate,
     plate_stack,
@@ -41,6 +43,7 @@
     "condition",
     "deterministic",
     "do",
+    "enable_module_local_param",
     "enable_validation",
     "factor",
     "get_param_store",
@@ -49,6 +52,7 @@
     "log",
     "markov",
     "module",
+    "module_local_param_enabled",
     "param",
     "plate",
     "plate",

pyro/infer/elbo.py

@@ -5,13 +5,26 @@
 import warnings
 from abc import ABCMeta, abstractmethod
 
+import torch
+
 import pyro
 import pyro.poutine as poutine
 from pyro.infer.util import is_validation_enabled
 from pyro.poutine.util import prune_subsample_sites
 from pyro.util import check_site_shape
 
 
+class ELBOModule(torch.nn.Module):
+    def __init__(self, model: torch.nn.Module, guide: torch.nn.Module, elbo: "ELBO"):
+        super().__init__()
+        self.model = model
+        self.guide = guide
+        self.elbo = elbo
+
+    def forward(self, *args, **kwargs):
+        return self.elbo.differentiable_loss(self.model, self.guide, *args, **kwargs)
+
+
 class ELBO(object, metaclass=ABCMeta):
     """
     :class:`ELBO` is the top-level interface for stochastic variational
@@ -23,6 +36,40 @@ class ELBO(object, metaclass=ABCMeta):
     :class:`~pyro.infer.tracegraph_elbo.TraceGraph_ELBO`, or
     :class:`~pyro.infer.traceenum_elbo.TraceEnum_ELBO`.
 
+    .. note:: Derived classes now provide a more idiomatic PyTorch interface via
+        :meth:`__call__` for (model, guide) pairs that are :class:`~torch.nn.Module` s,
+        which is useful for integrating Pyro's variational inference tooling with
+        standard PyTorch interfaces like :class:`~torch.optim.Optimizer` s
+        and the large ecosystem of libraries like PyTorch Lightning
+        and the PyTorch JIT that work with these interfaces::
+
+            model = Model()
+            guide = pyro.infer.autoguide.AutoNormal(model)
+
+            elbo_ = pyro.infer.Trace_ELBO(num_particles=10)
+
+            # Fix the model/guide pair
+            elbo = elbo_(model, guide)
+
+            # perform any data-dependent initialization
+            elbo(data)
+
+            optim = torch.optim.Adam(elbo.parameters(), lr=0.001)
+
+            for _ in range(100):
+                optim.zero_grad()
+                loss = elbo(data)
+                loss.backward()
+                optim.step()
+
+        Note that Pyro's global parameter store may cause this new interface to
+        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.
+
     :param num_particles: The number of particles/samples used to form the ELBO
         (gradient) estimators.
     :param int max_plate_nesting: Optional bound on max number of nested
@@ -86,6 +133,13 @@ def __init__(
         self.jit_options = jit_options
         self.tail_adaptive_beta = tail_adaptive_beta
 
+    def __call__(self, model: torch.nn.Module, guide: torch.nn.Module) -> ELBOModule:
+        """
+        Given a model and guide, returns a :class:`~torch.nn.Module` which
+        computes the ELBO loss when called with arguments to the model and guide.
+        """
+        return ELBOModule(model, guide, self)
+
     def _guess_max_plate_nesting(self, model, guide, args, kwargs):
         """
         Guesses max_plate_nesting by running the (model,guide) pair once

pyro/nn/module.py

@@ -169,6 +169,10 @@ def _unconstrain(constrained_value, constraint):
         return torch.nn.Parameter(unconstrained_value)
 
 
+def _is_module_local_param_enabled() -> bool:
+    return pyro.poutine.runtime._PYRO_MODULE_LOCAL_PARAM
+
+
 class _Context:
     """
     Sometimes-active cache for ``PyroModule.__call__()`` contexts.
@@ -178,15 +182,23 @@ def __init__(self):
         self.active = 0
         self.cache = {}
         self.used = False
+        if _is_module_local_param_enabled():
+            self.param_state = {"params": {}, "constraints": {}}
 
     def __enter__(self):
+        if not self.active and _is_module_local_param_enabled():
+            self._param_ctx = pyro.get_param_store().scope(state=self.param_state)
+            self.param_state = self._param_ctx.__enter__()
         self.active += 1
         self.used = True
 
     def __exit__(self, type, value, traceback):
         self.active -= 1
         if not self.active:
             self.cache.clear()
+            if _is_module_local_param_enabled():
+                self._param_ctx.__exit__(type, value, traceback)
+                del self._param_ctx
 
     def get(self, name):
         if self.active:
@@ -409,6 +421,22 @@ def named_pyro_params(self, prefix="", recurse=True):
             yield elem
 
     def _pyro_set_supermodule(self, name, context):
+        if _is_module_local_param_enabled():
+            with pyro.get_param_store().scope(
+                state=self._pyro_context.param_state
+            ) as vanilla_param_state:
+                # make sure we only propagate vanilla pyro.param data
+                #   in the _Context param_state - all others are handled
+                #   by the standard PyroModule/torch.nn.Module semantics
+                for param_name in list(pyro.get_param_store().keys()):
+                    # hack to detect vanilla pyro.params:
+                    #   vanilla pyro.params unattached to any Module
+                    #   should not have a pyromodule prefix
+                    if not param_name.startswith(self._pyro_name) or isinstance(
+                        getattr(self, param_name, None), torch.nn.Parameter
+                    ):
+                        del pyro.get_param_store()[param_name]
+            context.param_state.update(vanilla_param_state)
         self._pyro_name = name
         self._pyro_context = context
         for key, value in self._modules.items():

pyro/poutine/runtime.py

@@ -15,6 +15,9 @@
 # the global ParamStore
 _PYRO_PARAM_STORE = ParamStoreDict()
 
+# toggle usage of local param stores in PyroModules
+_PYRO_MODULE_LOCAL_PARAM = False
+
 
 class _DimAllocator:
     """

pyro/primitives.py

@@ -562,3 +562,35 @@ def validation_enabled(is_validate=True):
         dist.enable_validation(distribution_validation_status)
         infer.enable_validation(infer_validation_status)
         poutine.enable_validation(poutine_validation_status)
+
+
+def enable_module_local_param(is_enabled: bool = False) -> None:
+    """
+    Toggles the behavior of :class:`~pyro.nn.module.PyroModule` to use
+    local parameters instead of global parameters.
+
+    When this feature is enabled, :class:`~pyro.nn.module.PyroModule`
+    instances will not share parameters with other instances of the same
+    class through Pyro's global parameter store. Instead, each instance
+    will have its own local parameters, just like a standard :class:`torch.nn.Module`.
+
+    .. note:: This feature is disabled by default to ensure backwards compatibility
+        of :class:`~pyro.nn.module.PyroModule` with existing Pyro code.
+
+    :param bool is_enabled: (optional; defaults to False) whether to
+        enable local parameters.
+    """
+    poutine.runtime._PYRO_MODULE_LOCAL_PARAM = is_enabled
+
+
+@contextmanager
+def module_local_param_enabled(is_enabled=False):
+    """
+    Context manager to temporarily toggle local parameter stores in PyroModules.
+    """
+    old_flag = poutine.runtime._PYRO_MODULE_LOCAL_PARAM
+    poutine.runtime._PYRO_MODULE_LOCAL_PARAM = is_enabled
+    try:
+        yield
+    finally:
+        poutine.runtime._PYRO_MODULE_LOCAL_PARAM = old_flag

tests/conftest.py

@@ -34,6 +34,7 @@ def pytest_runtest_setup(item):
     if test_initialize_marker:
         rng_seed = test_initialize_marker.kwargs["rng_seed"]
         pyro.set_rng_seed(rng_seed)
+    pyro.enable_module_local_param(False)
 
 
 def pytest_addoption(parser):