Repository: pyro-ppl/pyro

Issue:
[discussion] Multiple guides/models in the same process
Following on from [this brief discussion](https://forum.pyro.ai/t/multiple-guides-models-in-the-same-process/3240) on the forum...

I would like to compare the results of using different guides for the same model (or the results from different models with different guides). Similar to this post ([Multiple guides for same model ](https://forum.pyro.ai/t/multiple-guides-for-same-model/1316)) it is of course possible to clear the global parameter store between runs of different models/guides, but that makes interactive comparison of the models/guides rather cumbersome. I’ve been using the following context manager to create multiple parameter stores (i.e. within the scope, a particular “local” parameter store is used rather than the “global” one).

Is this a sensible approach or is there something more straightforward that I’ve been missing? Is there interest in integrating this type of behaviour into the pyro codebase? This may be related to #502 in terms of global parameter stores sometimes being challenging.

```python
import pyro
from pyro.params.param_store import ParamStoreDict


class paramstore_scope:
    """
    Context manager for using multiple parameter stores within the same process.
    """
    GLOBAL_STORE = None
    
    def __init__(self, paramstore=None):
        self.paramstore = paramstore or ParamStoreDict()
        
    def __enter__(self):
        if self.GLOBAL_STORE is not None:
            raise RuntimeError('paramstore scopes cannot be nested')
        # Store the global paramstore.
        self.__class__.GLOBAL_STORE = pyro.poutine.runtime._PYRO_PARAM_STORE
        self._setup_paramstore(self.paramstore)
        return self
        
    def __exit__(self, *args):
        # Restore the global paramstore.
        self._setup_paramstore(self.GLOBAL_STORE)
        self.__class__.GLOBAL_STORE = None
        
    def _setup_paramstore(self, paramstore):
        # Replace all imports of the global param store with our version.
        pyro.poutine.runtime._PYRO_PARAM_STORE = \
            pyro.primitives._PYRO_PARAM_STORE = \
            pyro.nn.module._PYRO_PARAM_STORE = \
            paramstore
        # Replace static use of the imported param store.
        pyro.primitives._param = \
            pyro.poutine.runtime.effectful(paramstore.get_param, type="param")
```

An example looks like this.

```python
with paramstore_scope() as scope1:
    # Do your inference for model 1 (which will leave the global paramstore unchanged).

with paramstore_scope() as scope2:
    # Do your inference for model 2 (which will leave the global paramstore unchanged).

# Interactively explore the different models.
```


Discussion:

Resolving patch:
diff --git a/pyro/params/param_store.py b/pyro/params/param_store.py
--- a/pyro/params/param_store.py
+++ b/pyro/params/param_store.py
@@ -4,6 +4,7 @@
 import re
 import warnings
 import weakref
+from contextlib import contextmanager
 
 import torch
 from torch.distributions import constraints, transform_to
@@ -235,19 +236,22 @@ def param_name(self, p):
         """
         return self._param_to_name.get(p)
 
-    def get_state(self):
+    # -------------------------------------------------------------------------------
+    # Persistence interface
+
+    def get_state(self) -> dict:
         """
         Get the ParamStore state.
         """
         state = {
-            "params": self._params,
-            "constraints": self._constraints,
+            "params": self._params.copy(),
+            "constraints": self._constraints.copy(),
         }
         return state
 
-    def set_state(self, state):
+    def set_state(self, state: dict):
         """
-        Set the ParamStore state using state from a previous get_state() call
+        Set the ParamStore state using state from a previous :meth:`get_state` call
         """
         assert isinstance(state, dict), "malformed ParamStore state"
         assert set(state.keys()) == set(
@@ -266,7 +270,7 @@ def set_state(self, state):
 
     def save(self, filename):
         """
-        Save parameters to disk
+        Save parameters to file
 
         :param filename: file name to save to
         :type filename: str
@@ -276,7 +280,7 @@ def save(self, filename):
 
     def load(self, filename, map_location=None):
         """
-        Loads parameters from disk
+        Loads parameters from file
 
         .. note::
 
@@ -295,6 +299,44 @@ def load(self, filename, map_location=None):
             state = torch.load(input_file, map_location)
         self.set_state(state)
 
+    @contextmanager
+    def scope(self, state=None) -> dict:
+        """
+        Context manager for using multiple parameter stores within the same process.
+
+        This is a thin wrapper around :meth:`get_state`, :meth:`clear`, and
+        :meth:`set_state`. For large models where memory space is limiting, you
+        may want to instead manually use :meth:`save`, :meth:`clear`, and
+        :meth:`load`.
+
+        Example usage::
+
+            param_store = pyro.get_param_store()
+
+            # Train multiple models, while avoiding param name conflicts.
+            with param_store.scope() as scope1:
+                # ...Train one model,guide pair...
+            with param_store.scope() as scope2:
+                # ...Train another model,guide pair...
+
+            # Now evaluate each, still avoiding name conflicts.
+            with param_store.scope(scope1):  # loads the first model's scope
+               # ...evaluate the first model...
+            with param_store.scope(scope2):  # loads the second model's scope
+               # ...evaluate the second model...
+        """
+        if state is None:
+            state = {"params": {}, "constraints": {}}
+        old_state = self.get_state()
+        try:
+            self.clear()
+            self.set_state(state)
+            yield state
+            state.update(self.get_state())
+        finally:
+            self.clear()
+            self.set_state(old_state)
+
 
 # used to create fully-formed param names, e.g. mymodule$$$mysubmodule.weight
 _MODULE_NAMESPACE_DIVIDER = "$$$"


Test patch:
diff --git a/tests/params/test_param.py b/tests/params/test_param.py
--- a/tests/params/test_param.py
+++ b/tests/params/test_param.py
@@ -149,3 +149,97 @@ def test_dict_interface():
     assert list(param_store.keys()) == []
     assert list(key for key, value in param_store.items()) == []
     assert len(list(param_store.values())) == 0
+
+
+def test_scope():
+    x0 = torch.randn(())
+    z0 = torch.randn(5).exp()
+    x1 = torch.randn(3)
+    y1 = torch.randn(2, 1).exp()
+    y2 = torch.randn(2, 1).exp()
+    z2 = torch.randn(1, 4).exp()
+    z2 /= z2.sum()
+
+    constraint_table = {
+        "z0": constraints.positive,
+        "y1": constraints.positive,
+        "y2": constraints.positive,
+        "z2": constraints.simplex,
+    }
+
+    def check_constraint(name):
+        expected = constraint_table.get(name, constraints.real)
+        actual = param_store._constraints[name[:1]]
+        assert actual == expected
+
+    param_store = pyro.get_param_store()
+
+    # Create base scope.
+    assert not param_store
+    pyro.param("x", x0)
+    pyro.param("z", z0, constraint=constraints.positive)
+    assert set(param_store) == {"x", "z"}
+    assert_equal(pyro.param("x"), x0)
+    assert_equal(pyro.param("z"), z0)
+    check_constraint("x0")
+    check_constraint("z0")
+
+    # Create scope1.
+    with param_store.scope() as scope1:
+        assert not param_store
+        pyro.param("x", x1)
+        pyro.param("y", y1, constraint=constraints.positive)
+        assert set(param_store) == {"x", "y"}
+        assert_equal(pyro.param("x"), x1)
+        assert_equal(pyro.param("y"), y1)
+        check_constraint("x1")
+        check_constraint("y1")
+    assert set(param_store) == {"x", "z"}
+    assert_equal(pyro.param("x"), x0)
+    assert_equal(pyro.param("z"), z0)
+    check_constraint("x0")
+    check_constraint("z0")
+
+    # Create scope2.
+    with param_store.scope() as scope2:
+        assert not param_store
+        pyro.param("y", y2, constraint=constraints.positive)
+        pyro.param("z", z2, constraint=constraints.simplex)
+        assert set(param_store) == {"y", "z"}
+        assert_equal(pyro.param("y"), y2)
+        assert_equal(pyro.param("z"), z2)
+        check_constraint("y2")
+        check_constraint("z2")
+    assert set(param_store) == {"x", "z"}
+    assert_equal(pyro.param("x"), x0)
+    assert_equal(pyro.param("z"), z0)
+    check_constraint("x0")
+    check_constraint("z0")
+
+    # Reload scope1.
+    with param_store.scope(scope1) as s:
+        assert s is scope1
+        assert set(param_store) == {"x", "y"}
+        assert_equal(pyro.param("x"), x1)
+        assert_equal(pyro.param("y"), y1)
+        check_constraint("x1")
+        check_constraint("y1")
+    assert set(param_store) == {"x", "z"}
+    assert_equal(pyro.param("x"), x0)
+    assert_equal(pyro.param("z"), z0)
+    check_constraint("x0")
+    check_constraint("z0")
+
+    # Reload scope2.
+    with param_store.scope(scope2) as s:
+        assert s is scope2
+        assert set(param_store) == {"y", "z"}
+        assert_equal(pyro.param("y"), y2)
+        assert_equal(pyro.param("z"), z2)
+        check_constraint("y2")
+        check_constraint("z2")
+    assert set(param_store) == {"x", "z"}
+    assert_equal(pyro.param("x"), x0)
+    assert_equal(pyro.param("z"), z0)
+    check_constraint("x0")
+    check_constraint("z0")