From 06a91220cbe3e686f0b0a95776109d4ca42a90dd Mon Sep 17 00:00:00 2001 From: Eli Date: Wed, 19 Oct 2022 13:10:59 -0400 Subject: [PATCH 01/17] Much simpler local params in PyroModule --- pyro/nn/module.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyro/nn/module.py b/pyro/nn/module.py index 7a87631e2e..3088e0456a 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -180,6 +180,9 @@ def __init__(self): self.used = False def __enter__(self): + if not self.active: + self._param_ctx = pyro.get_param_store().scope(state=None) + self._param_ctx.__enter__() self.active += 1 self.used = True @@ -187,6 +190,8 @@ def __exit__(self, type, value, traceback): self.active -= 1 if not self.active: self.cache.clear() + self._param_ctx.__exit__(type, value, traceback) + del self._param_ctx def get(self, name): if self.active: From fb5d282a3503c1237b3d3672a2258a477256062a Mon Sep 17 00:00:00 2001 From: Eli Date: Wed, 19 Oct 2022 13:18:21 -0400 Subject: [PATCH 02/17] add toggle --- pyro/infer/elbo.py | 16 ++++++++++++++++ pyro/nn/module.py | 9 +++++---- pyro/poutine/runtime.py | 12 ++++++++++++ 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/pyro/infer/elbo.py b/pyro/infer/elbo.py index 3abe07d748..991d278350 100644 --- a/pyro/infer/elbo.py +++ b/pyro/infer/elbo.py @@ -5,6 +5,8 @@ import warnings from abc import ABCMeta, abstractmethod +import torch + import pyro import pyro.poutine as poutine from pyro.infer.util import is_validation_enabled @@ -12,6 +14,17 @@ from pyro.util import check_site_shape +class _ELBOModule(torch.nn.Module): + def __init__(self, model, guide, 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 @@ -86,6 +99,9 @@ def __init__( self.jit_options = jit_options self.tail_adaptive_beta = tail_adaptive_beta + def __call__(self, model, 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 diff --git a/pyro/nn/module.py b/pyro/nn/module.py index 3088e0456a..ea8bf04600 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -21,7 +21,7 @@ import pyro from pyro.ops.provenance import detach_provenance -from pyro.poutine.runtime import _PYRO_PARAM_STORE +from pyro.poutine.runtime import _PYRO_PARAM_STORE, _module_local_param_enabled class PyroParam(namedtuple("PyroParam", ("init_value", "constraint", "event_dim"))): @@ -180,7 +180,7 @@ def __init__(self): self.used = False def __enter__(self): - if not self.active: + if not self.active and _module_local_param_enabled(): self._param_ctx = pyro.get_param_store().scope(state=None) self._param_ctx.__enter__() self.active += 1 @@ -190,8 +190,9 @@ def __exit__(self, type, value, traceback): self.active -= 1 if not self.active: self.cache.clear() - self._param_ctx.__exit__(type, value, traceback) - del self._param_ctx + if _module_local_param_enabled(): + self._param_ctx.__exit__(type, value, traceback) + del self._param_ctx def get(self, name): if self.active: diff --git a/pyro/poutine/runtime.py b/pyro/poutine/runtime.py index 59b27c8911..612d51fd9a 100644 --- a/pyro/poutine/runtime.py +++ b/pyro/poutine/runtime.py @@ -15,6 +15,18 @@ # the global ParamStore _PYRO_PARAM_STORE = ParamStoreDict() +# toggle usage of local param stores in PyroModules +_PYRO_MODULE_LOCAL_PARAM = False + + +def enable_module_local_param(flag: bool) -> None: + global _PYRO_MODULE_LOCAL_PARAM + _PYRO_MODULE_LOCAL_PARAM = flag + + +def _module_local_param_enabled(): + return _PYRO_MODULE_LOCAL_PARAM + class _DimAllocator: """ From ac21b8d4eece56696a026683323aef728e243f8f Mon Sep 17 00:00:00 2001 From: Eli Date: Wed, 19 Oct 2022 13:21:19 -0400 Subject: [PATCH 03/17] import --- pyro/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyro/__init__.py b/pyro/__init__.py index c8c07e6eb2..6c25ea6086 100644 --- a/pyro/__init__.py +++ b/pyro/__init__.py @@ -5,6 +5,7 @@ from pyro.infer.inspect import render_model from pyro.logger import log from pyro.poutine import condition, do, markov +from pyro.poutine.runtime import enable_module_local_param from pyro.primitives import ( barrier, clear_param_store, @@ -41,6 +42,7 @@ "condition", "deterministic", "do", + "enable_module_local_param", "enable_validation", "factor", "get_param_store", From c789c3e782d3551f234edda7488bd71d5efa0ed7 Mon Sep 17 00:00:00 2001 From: Eli Date: Fri, 21 Oct 2022 17:05:19 -0400 Subject: [PATCH 04/17] move code around --- pyro/__init__.py | 4 +++- pyro/nn/module.py | 10 +++++++--- pyro/poutine/runtime.py | 9 --------- pyro/primitives.py | 32 ++++++++++++++++++++++++++++++++ 4 files changed, 42 insertions(+), 13 deletions(-) diff --git a/pyro/__init__.py b/pyro/__init__.py index 6c25ea6086..65d0d8f3fa 100644 --- a/pyro/__init__.py +++ b/pyro/__init__.py @@ -5,17 +5,18 @@ from pyro.infer.inspect import render_model from pyro.logger import log from pyro.poutine import condition, do, markov -from pyro.poutine.runtime import enable_module_local_param from pyro.primitives import ( 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, @@ -51,6 +52,7 @@ "log", "markov", "module", + "module_local_param_enabled", "param", "plate", "plate", diff --git a/pyro/nn/module.py b/pyro/nn/module.py index ea8bf04600..cae6a1e8d2 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -21,7 +21,7 @@ import pyro from pyro.ops.provenance import detach_provenance -from pyro.poutine.runtime import _PYRO_PARAM_STORE, _module_local_param_enabled +from pyro.poutine.runtime import _PYRO_PARAM_STORE class PyroParam(namedtuple("PyroParam", ("init_value", "constraint", "event_dim"))): @@ -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. @@ -180,7 +184,7 @@ def __init__(self): self.used = False def __enter__(self): - if not self.active and _module_local_param_enabled(): + if not self.active and _is_module_local_param_enabled(): self._param_ctx = pyro.get_param_store().scope(state=None) self._param_ctx.__enter__() self.active += 1 @@ -190,7 +194,7 @@ def __exit__(self, type, value, traceback): self.active -= 1 if not self.active: self.cache.clear() - if _module_local_param_enabled(): + if _is_module_local_param_enabled(): self._param_ctx.__exit__(type, value, traceback) del self._param_ctx diff --git a/pyro/poutine/runtime.py b/pyro/poutine/runtime.py index 612d51fd9a..85c6e19274 100644 --- a/pyro/poutine/runtime.py +++ b/pyro/poutine/runtime.py @@ -19,15 +19,6 @@ _PYRO_MODULE_LOCAL_PARAM = False -def enable_module_local_param(flag: bool) -> None: - global _PYRO_MODULE_LOCAL_PARAM - _PYRO_MODULE_LOCAL_PARAM = flag - - -def _module_local_param_enabled(): - return _PYRO_MODULE_LOCAL_PARAM - - class _DimAllocator: """ Dimension allocator for internal use by :class:`plate`. diff --git a/pyro/primitives.py b/pyro/primitives.py index ca67d7b061..0fa34e3ca0 100644 --- a/pyro/primitives.py +++ b/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 From 9835ee954f069804849c207ea8aaa1349e398551 Mon Sep 17 00:00:00 2001 From: Eli Date: Sat, 22 Oct 2022 14:09:29 -0400 Subject: [PATCH 05/17] add tests --- tests/nn/test_module.py | 232 ++++++++++++++++++++++++---------------- 1 file changed, 142 insertions(+), 90 deletions(-) diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 41198b9f2a..1de079031c 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -66,7 +66,8 @@ def forward(self, *args, **kwargs): svi.step(data) -def test_names(): +@pytest.mark.parametrize("local_params", [True, False]) +def test_names(local_params): class Model(PyroModule): def __init__(self): super().__init__() @@ -86,34 +87,39 @@ def forward(self): self.p.v self.p.w - model = Model() - - # Check named_parameters. - expected = { - "x", - "y_unconstrained", - "m.u", - "p.v", - "p.w_unconstrained", - } - actual = set(name for name, _ in model.named_parameters()) - assert actual == expected - - # Check pyro.param names. - expected = {"x", "y", "m$$$u", "p.v", "p.w"} - with poutine.trace(param_only=True) as param_capture: - model() - actual = { - name - for name, site in param_capture.trace.nodes.items() - if site["type"] == "param" - } - assert actual == expected - - # Check pyro_parameters method - expected = {"x", "y", "m.u", "p.v", "p.w"} - actual = set(k for k, v in model.named_pyro_params()) - assert actual == expected + with pyro.module_local_param_enabled(local_params): + model = Model() + + # Check named_parameters. + expected = { + "x", + "y_unconstrained", + "m.u", + "p.v", + "p.w_unconstrained", + } + actual = set(name for name, _ in model.named_parameters()) + assert actual == expected + + # Check pyro.param names. + expected = {"x", "y", "m$$$u", "p.v", "p.w"} + with poutine.trace(param_only=True) as param_capture: + model() + actual = { + name + for name, site in param_capture.trace.nodes.items() + if site["type"] == "param" + } + assert actual == expected + if local_params: + assert set(pyro.get_param_store().keys()) == set() + else: + assert set(pyro.get_param_store().keys()) == expected + + # Check pyro_parameters method + expected = {"x", "y", "m.u", "p.v", "p.w"} + actual = set(k for k, v in model.named_pyro_params()) + assert actual == expected def test_delete(): @@ -258,7 +264,8 @@ def test_constraints(shape, constraint_): assert not hasattr(module, "x_unconstrained") -def test_clear(): +@pytest.mark.parametrize("local_params", [True, False]) +def test_clear(local_params): class Model(PyroModule): def __init__(self): super().__init__() @@ -272,28 +279,33 @@ def __init__(self): def forward(self): return [x.clone() for x in [self.x, self.m.weight, self.m.bias, self.p.x]] - assert set(pyro.get_param_store().keys()) == set() - m = Model() - state0 = m() - assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} - - # mutate - for x in pyro.get_param_store().values(): - x.unconstrained().data += torch.randn(()) - state1 = m() - for x, y in zip(state0, state1): - assert not (x == y).all() - assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} - - clear(m) - del m - assert set(pyro.get_param_store().keys()) == set() - - m = Model() - state2 = m() - assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} - for actual, expected in zip(state2, state0): - assert_equal(actual, expected) + with pyro.module_local_param_enabled(local_params): + m = Model() + state0 = m() + + # mutate + for _, x in m.named_pyro_params(): + x.unconstrained().data += torch.randn(()) + state1 = m() + for x, y in zip(state0, state1): + assert not (x == y).all() + + if local_params: + assert set(pyro.get_param_store().keys()) == set() + else: + assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} + clear(m) + del m + assert set(pyro.get_param_store().keys()) == set() + + m = Model() + state2 = m() + if local_params: + assert set(pyro.get_param_store().keys()) == set() + else: + assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} + for actual, expected in zip(state2, state0): + assert_equal(actual, expected) def test_sample(): @@ -411,23 +423,25 @@ def forward(self): @pytest.mark.parametrize("Model", [AttributeModel, DecoratorModel]) @pytest.mark.parametrize("size", [1, 2]) -def test_decorator(Model, size): - model = Model(size) - for i in range(2): - trace = poutine.trace(model).get_trace() - assert set(trace.nodes.keys()) == {"_INPUT", "x", "y", "z", "s", "t", "_RETURN"} - - assert trace.nodes["x"]["type"] == "param" - assert trace.nodes["y"]["type"] == "param" - assert trace.nodes["z"]["type"] == "param" - assert trace.nodes["s"]["type"] == "sample" - assert trace.nodes["t"]["type"] == "sample" - - assert trace.nodes["x"]["value"].shape == (size,) - assert trace.nodes["y"]["value"].shape == (size,) - assert trace.nodes["z"]["value"].shape == (size,) - assert trace.nodes["s"]["value"].shape == () - assert trace.nodes["t"]["value"].shape == (size,) +@pytest.mark.parametrize("local_params", [False, True]) +def test_decorator(Model, size, local_params): + with pyro.module_local_param_enabled(local_params): + model = Model(size) + for i in range(2): + trace = poutine.trace(model).get_trace() + assert set(trace.nodes.keys()) == {"_INPUT", "x", "y", "z", "s", "t", "_RETURN"} + + assert trace.nodes["x"]["type"] == "param" + assert trace.nodes["y"]["type"] == "param" + assert trace.nodes["z"]["type"] == "param" + assert trace.nodes["s"]["type"] == "sample" + assert trace.nodes["t"]["type"] == "sample" + + assert trace.nodes["x"]["value"].shape == (size,) + assert trace.nodes["y"]["value"].shape == (size,) + assert trace.nodes["z"]["value"].shape == (size,) + assert trace.nodes["s"]["value"].shape == () + assert trace.nodes["t"]["value"].shape == (size,) def test_mixin_factory(): @@ -553,27 +567,42 @@ def test_torch_serialize_attributes(): assert actual_names == expected_names -def test_torch_serialize_decorators(): - module = DecoratorModel(3) - module() # initialize - - # Work around https://github.com/pytorch/pytorch/issues/27972 - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=UserWarning) - f = io.BytesIO() - torch.save(module, f) - pyro.clear_param_store() - f.seek(0) - actual = torch.load(f) - - assert_equal(actual.x, module.x) - assert_equal(actual.y, module.y) - assert_equal(actual.z, module.z) - assert actual.s.shape == module.s.shape - assert actual.t.shape == module.t.shape - actual_names = {name for name, _ in actual.named_parameters()} - expected_names = {name for name, _ in module.named_parameters()} - assert actual_names == expected_names +@pytest.mark.parametrize("local_params", [True, False]) +def test_torch_serialize_decorators(local_params): + with pyro.module_local_param_enabled(local_params): + module = DecoratorModel(3) + module() # initialize + + module2 = DecoratorModel(3) + module2() # initialize + + # Work around https://github.com/pytorch/pytorch/issues/27972 + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=UserWarning) + f = io.BytesIO() + torch.save(module, f) + pyro.clear_param_store() + f.seek(0) + actual = torch.load(f) + + assert_equal(actual.x, module.x) + assert_equal(actual.y, module.y) + assert_equal(actual.z, module.z) + assert actual.s.shape == module.s.shape + assert actual.t.shape == module.t.shape + actual_names = {name for name, _ in actual.named_parameters()} + expected_names = {name for name, _ in module.named_parameters()} + assert actual_names == expected_names + + actual() + if local_params: + assert len(set(pyro.get_param_store().keys())) == 0 + assert torch.all(module.y != module2.y) + assert torch.all(actual.y != module2.y) + else: + assert len(set(pyro.get_param_store().keys())) > 0 + assert_equal(module.y, module2.y) + assert_equal(actual.y, module2.y) def test_pyro_serialize(): @@ -636,3 +665,26 @@ def test_bayesian_gru(): assert output.shape == (seq_len, batch_size, hidden_size) output2, _ = gru(input_) assert not torch.allclose(output2, output) + + +@pytest.mark.parametrize("local_params", [True, False]) +def test_pyro_module_mixin_local_params(local_params): + with pyro.module_local_param_enabled(local_params): + input_size = 2 + hidden_size = 3 + batch_size = 4 + input_ = torch.randn(batch_size, input_size) + + c = PyroModule[nn.Linear] + m = c(input_size, hidden_size) + output = m(input_) + + del m + + m2 = c(input_size, hidden_size) + output2 = m2(input_) + + if local_params: + assert not torch.allclose(output2, output) + else: + assert torch.allclose(output2, output) From 2cb054cba385f1da6799e3aa5ee1cc106fdcd1db Mon Sep 17 00:00:00 2001 From: Eli Date: Sun, 23 Oct 2022 16:35:16 -0400 Subject: [PATCH 06/17] remove weird test --- tests/nn/test_module.py | 120 ++++++++++++++++++++-------------------- 1 file changed, 59 insertions(+), 61 deletions(-) diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 1de079031c..612ec193f3 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -293,7 +293,12 @@ def forward(self): if local_params: assert set(pyro.get_param_store().keys()) == set() else: - assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} + assert set(pyro.get_param_store().keys()) == { + "x", + "m$$$weight", + "m$$$bias", + "p.x", + } clear(m) del m assert set(pyro.get_param_store().keys()) == set() @@ -303,7 +308,12 @@ def forward(self): if local_params: assert set(pyro.get_param_store().keys()) == set() else: - assert set(pyro.get_param_store().keys()) == {"x", "m$$$weight", "m$$$bias", "p.x"} + assert set(pyro.get_param_store().keys()) == { + "x", + "m$$$weight", + "m$$$bias", + "p.x", + } for actual, expected in zip(state2, state0): assert_equal(actual, expected) @@ -339,42 +349,45 @@ def __init__(self, in_features, out_features): svi.step(data) -def test_cache(): - class MyModule(PyroModule): - def forward(self): - return [self.gather(), self.gather()] - - def gather(self): - return { - "a": self.a, - "b": self.b, - "c": self.c, - "p.d": self.p.d, - "p.e": self.p.e, - "p.f": self.p.f, - } - - module = MyModule() - module.a = nn.Parameter(torch.tensor(0.0)) - module.b = PyroParam(torch.tensor(1.0), constraint=constraints.positive) - module.c = PyroSample(dist.Normal(0, 1)) - module.p = PyroModule() - module.p.d = nn.Parameter(torch.tensor(3.0)) - module.p.e = PyroParam(torch.tensor(4.0), constraint=constraints.positive) - module.p.f = PyroSample(dist.Normal(0, 1)) - - assert module._pyro_context is module.p._pyro_context - - # Check that results are cached with an invocation of .__call__(). - result1 = module() - actual, expected = result1 - for key in ["a", "c", "p.d", "p.f"]: - assert actual[key] is expected[key], key +@pytest.mark.parametrize("local_params", [True, False]) +def test_cache(local_params): + with pyro.module_local_param_enabled(local_params): - # Check that results are not cached across invocations of .__call__(). - result2 = module() - for key in ["b", "c", "p.e", "p.f"]: - assert result1[0] is not result2[0], key + class MyModule(PyroModule): + def forward(self): + return [self.gather(), self.gather()] + + def gather(self): + return { + "a": self.a, + "b": self.b, + "c": self.c, + "p.d": self.p.d, + "p.e": self.p.e, + "p.f": self.p.f, + } + + module = MyModule() + module.a = nn.Parameter(torch.tensor(0.0)) + module.b = PyroParam(torch.tensor(1.0), constraint=constraints.positive) + module.c = PyroSample(dist.Normal(0, 1)) + module.p = PyroModule() + module.p.d = nn.Parameter(torch.tensor(3.0)) + module.p.e = PyroParam(torch.tensor(4.0), constraint=constraints.positive) + module.p.f = PyroSample(dist.Normal(0, 1)) + + assert module._pyro_context is module.p._pyro_context + + # Check that results are cached with an invocation of .__call__(). + result1 = module() + actual, expected = result1 + for key in ["a", "c", "p.d", "p.f"]: + assert actual[key] is expected[key], key + + # Check that results are not cached across invocations of .__call__(). + result2 = module() + for key in ["b", "c", "p.e", "p.f"]: + assert result1[0] is not result2[0], key class AttributeModel(PyroModule): @@ -429,7 +442,15 @@ def test_decorator(Model, size, local_params): model = Model(size) for i in range(2): trace = poutine.trace(model).get_trace() - assert set(trace.nodes.keys()) == {"_INPUT", "x", "y", "z", "s", "t", "_RETURN"} + assert set(trace.nodes.keys()) == { + "_INPUT", + "x", + "y", + "z", + "s", + "t", + "_RETURN", + } assert trace.nodes["x"]["type"] == "param" assert trace.nodes["y"]["type"] == "param" @@ -665,26 +686,3 @@ def test_bayesian_gru(): assert output.shape == (seq_len, batch_size, hidden_size) output2, _ = gru(input_) assert not torch.allclose(output2, output) - - -@pytest.mark.parametrize("local_params", [True, False]) -def test_pyro_module_mixin_local_params(local_params): - with pyro.module_local_param_enabled(local_params): - input_size = 2 - hidden_size = 3 - batch_size = 4 - input_ = torch.randn(batch_size, input_size) - - c = PyroModule[nn.Linear] - m = c(input_size, hidden_size) - output = m(input_) - - del m - - m2 = c(input_size, hidden_size) - output2 = m2(input_) - - if local_params: - assert not torch.allclose(output2, output) - else: - assert torch.allclose(output2, output) From 755ae4aeccf1bd72c32e1592098b613bbac62cc4 Mon Sep 17 00:00:00 2001 From: Eli Date: Sun, 23 Oct 2022 17:05:45 -0400 Subject: [PATCH 07/17] add a test exercising the ElboModule interface --- tests/nn/test_module.py | 95 ++++++++++++++++++++++++++++++++++------- 1 file changed, 80 insertions(+), 15 deletions(-) diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 612ec193f3..def5322c3f 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -66,6 +66,72 @@ def forward(self, *args, **kwargs): svi.step(data) +@pytest.mark.parametrize("local_params", [True, False]) +@pytest.mark.parametrize("num_particles", [1, 2]) +@pytest.mark.parametrize("vectorize_particles", [True, False]) +@pytest.mark.parametrize("Autoguide", [pyro.infer.autoguide.AutoNormal]) +def test_svi_elbomodule_interface( + local_params, num_particles, vectorize_particles, Autoguide +): + class Model(PyroModule): + def __init__(self): + super().__init__() + self.loc = nn.Parameter(torch.zeros(2)) + self.scale = PyroParam(torch.ones(2), constraint=constraints.positive) + self.z = PyroSample( + lambda self: dist.Normal(self.loc, self.scale).to_event(1) + ) + + def forward(self, data): + loc, log_scale = self.z.unbind(-1) + with pyro.plate("data"): + pyro.sample("obs", dist.Cauchy(loc, log_scale.exp()), obs=data) + + with pyro.module_local_param_enabled(local_params): + data = torch.randn(5) + model = Model() + model(data) # initialize + + guide = Autoguide(model) + guide(data) # initialize + + elbo = Trace_ELBO( + vectorize_particles=vectorize_particles, num_particles=num_particles + ) + elbo = elbo(model, guide) + assert isinstance(elbo, torch.nn.Module) + assert set( + k[: -len("_unconstrained")] if k.endswith("_unconstrained") else k + for k, v in elbo.named_parameters() + ) == set("model." + k for k, v in model.named_pyro_params()) | set( + "guide." + k for k, v in guide.named_pyro_params() + ) + + adam = torch.optim.Adam(elbo.parameters(), lr=0.0001) + for _ in range(3): + adam.zero_grad() + loss = elbo(data) + loss.backward() + adam.step() + + guide2 = Autoguide(model) + guide2(data) # initialize + if local_params: + assert set(pyro.get_param_store().keys()) == set() + for (name, p), (name2, p2) in zip( + guide.named_parameters(), guide2.named_parameters() + ): + assert name == name2 + assert not torch.allclose(p, p2) + else: + assert set(pyro.get_param_store().keys()) != set() + for (name, p), (name2, p2) in zip( + guide.named_parameters(), guide2.named_parameters() + ): + assert name == name2 + assert torch.allclose(p, p2) + + @pytest.mark.parametrize("local_params", [True, False]) def test_names(local_params): class Model(PyroModule): @@ -351,22 +417,21 @@ def __init__(self, in_features, out_features): @pytest.mark.parametrize("local_params", [True, False]) def test_cache(local_params): - with pyro.module_local_param_enabled(local_params): - - class MyModule(PyroModule): - def forward(self): - return [self.gather(), self.gather()] - - def gather(self): - return { - "a": self.a, - "b": self.b, - "c": self.c, - "p.d": self.p.d, - "p.e": self.p.e, - "p.f": self.p.f, - } + class MyModule(PyroModule): + def forward(self): + return [self.gather(), self.gather()] + + def gather(self): + return { + "a": self.a, + "b": self.b, + "c": self.c, + "p.d": self.p.d, + "p.e": self.p.e, + "p.f": self.p.f, + } + with pyro.module_local_param_enabled(local_params): module = MyModule() module.a = nn.Parameter(torch.tensor(0.0)) module.b = PyroParam(torch.tensor(1.0), constraint=constraints.positive) From 5f7cd9445b12d7b3cf2b4ad509581fe969f19960 Mon Sep 17 00:00:00 2001 From: Eli Date: Sun, 23 Oct 2022 17:12:43 -0400 Subject: [PATCH 08/17] add default in top level conftest.py --- tests/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/conftest.py b/tests/conftest.py index 4103356fe3..5ed178546c 100644 --- a/tests/conftest.py +++ b/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): From c2eb114f9fa246793b014c38e02a2f845b04192f Mon Sep 17 00:00:00 2001 From: Eli Date: Sun, 23 Oct 2022 17:16:20 -0400 Subject: [PATCH 09/17] remove unnecessary change to test --- tests/nn/test_module.py | 46 ++++++++++++++++++++--------------------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index def5322c3f..c88a4c3a61 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -415,8 +415,7 @@ def __init__(self, in_features, out_features): svi.step(data) -@pytest.mark.parametrize("local_params", [True, False]) -def test_cache(local_params): +def test_cache(): class MyModule(PyroModule): def forward(self): return [self.gather(), self.gather()] @@ -431,28 +430,27 @@ def gather(self): "p.f": self.p.f, } - with pyro.module_local_param_enabled(local_params): - module = MyModule() - module.a = nn.Parameter(torch.tensor(0.0)) - module.b = PyroParam(torch.tensor(1.0), constraint=constraints.positive) - module.c = PyroSample(dist.Normal(0, 1)) - module.p = PyroModule() - module.p.d = nn.Parameter(torch.tensor(3.0)) - module.p.e = PyroParam(torch.tensor(4.0), constraint=constraints.positive) - module.p.f = PyroSample(dist.Normal(0, 1)) - - assert module._pyro_context is module.p._pyro_context - - # Check that results are cached with an invocation of .__call__(). - result1 = module() - actual, expected = result1 - for key in ["a", "c", "p.d", "p.f"]: - assert actual[key] is expected[key], key - - # Check that results are not cached across invocations of .__call__(). - result2 = module() - for key in ["b", "c", "p.e", "p.f"]: - assert result1[0] is not result2[0], key + module = MyModule() + module.a = nn.Parameter(torch.tensor(0.0)) + module.b = PyroParam(torch.tensor(1.0), constraint=constraints.positive) + module.c = PyroSample(dist.Normal(0, 1)) + module.p = PyroModule() + module.p.d = nn.Parameter(torch.tensor(3.0)) + module.p.e = PyroParam(torch.tensor(4.0), constraint=constraints.positive) + module.p.f = PyroSample(dist.Normal(0, 1)) + + assert module._pyro_context is module.p._pyro_context + + # Check that results are cached with an invocation of .__call__(). + result1 = module() + actual, expected = result1 + for key in ["a", "c", "p.d", "p.f"]: + assert actual[key] is expected[key], key + + # Check that results are not cached across invocations of .__call__(). + result2 = module() + for key in ["b", "c", "p.e", "p.f"]: + assert result1[0] is not result2[0], key class AttributeModel(PyroModule): From a3a848c60572e60d8ddc7e78bae1a0a87dd88a89 Mon Sep 17 00:00:00 2001 From: Eli Date: Sun, 23 Oct 2022 17:19:41 -0400 Subject: [PATCH 10/17] remove another unnecessary new test --- tests/nn/test_module.py | 44 ++++++++++++++++------------------------- 1 file changed, 17 insertions(+), 27 deletions(-) diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index c88a4c3a61..c95f36dd1b 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -499,33 +499,23 @@ def forward(self): @pytest.mark.parametrize("Model", [AttributeModel, DecoratorModel]) @pytest.mark.parametrize("size", [1, 2]) -@pytest.mark.parametrize("local_params", [False, True]) -def test_decorator(Model, size, local_params): - with pyro.module_local_param_enabled(local_params): - model = Model(size) - for i in range(2): - trace = poutine.trace(model).get_trace() - assert set(trace.nodes.keys()) == { - "_INPUT", - "x", - "y", - "z", - "s", - "t", - "_RETURN", - } - - assert trace.nodes["x"]["type"] == "param" - assert trace.nodes["y"]["type"] == "param" - assert trace.nodes["z"]["type"] == "param" - assert trace.nodes["s"]["type"] == "sample" - assert trace.nodes["t"]["type"] == "sample" - - assert trace.nodes["x"]["value"].shape == (size,) - assert trace.nodes["y"]["value"].shape == (size,) - assert trace.nodes["z"]["value"].shape == (size,) - assert trace.nodes["s"]["value"].shape == () - assert trace.nodes["t"]["value"].shape == (size,) +def test_decorator(Model, size): + model = Model(size) + for i in range(2): + trace = poutine.trace(model).get_trace() + assert set(trace.nodes.keys()) == {"_INPUT", "x", "y", "z", "s", "t", "_RETURN"} + + assert trace.nodes["x"]["type"] == "param" + assert trace.nodes["y"]["type"] == "param" + assert trace.nodes["z"]["type"] == "param" + assert trace.nodes["s"]["type"] == "sample" + assert trace.nodes["t"]["type"] == "sample" + + assert trace.nodes["x"]["value"].shape == (size,) + assert trace.nodes["y"]["value"].shape == (size,) + assert trace.nodes["z"]["value"].shape == (size,) + assert trace.nodes["s"]["value"].shape == () + assert trace.nodes["t"]["value"].shape == (size,) def test_mixin_factory(): From 81b1191cd1779143e19338e207720c82fcba839e Mon Sep 17 00:00:00 2001 From: Eli Date: Sun, 23 Oct 2022 18:06:52 -0400 Subject: [PATCH 11/17] document new interface to elbo --- pyro/infer/elbo.py | 44 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/pyro/infer/elbo.py b/pyro/infer/elbo.py index 991d278350..54cae9e1c2 100644 --- a/pyro/infer/elbo.py +++ b/pyro/infer/elbo.py @@ -15,7 +15,7 @@ class _ELBOModule(torch.nn.Module): - def __init__(self, model, guide, elbo): + def __init__(self, model: torch.nn.Module, guide: torch.nn.Module, elbo: "ELBO"): super().__init__() self.model = model self.guide = guide @@ -36,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 @@ -99,7 +133,13 @@ def __init__( self.jit_options = jit_options self.tail_adaptive_beta = tail_adaptive_beta - def __call__(self, model, guide): + def __call__( + self, model: torch.nn.Module, guide: torch.nn.Module + ) -> torch.nn.Module: + """ + 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): From 2c03fa03c5988a89a21c486efd1efe6a7c9fca86 Mon Sep 17 00:00:00 2001 From: Eli Date: Mon, 24 Oct 2022 20:06:13 -0400 Subject: [PATCH 12/17] address comments --- pyro/infer/elbo.py | 8 ++--- pyro/nn/module.py | 8 +++-- tests/nn/test_module.py | 79 +++++++++++++++++++++++++++++++++++++++-- 3 files changed, 86 insertions(+), 9 deletions(-) diff --git a/pyro/infer/elbo.py b/pyro/infer/elbo.py index 54cae9e1c2..05ebd6f2a9 100644 --- a/pyro/infer/elbo.py +++ b/pyro/infer/elbo.py @@ -14,7 +14,7 @@ from pyro.util import check_site_shape -class _ELBOModule(torch.nn.Module): +class ELBOModule(torch.nn.Module): def __init__(self, model: torch.nn.Module, guide: torch.nn.Module, elbo: "ELBO"): super().__init__() self.model = model @@ -133,14 +133,12 @@ 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 - ) -> torch.nn.Module: + 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) + return ELBOModule(model, guide, self) def _guess_max_plate_nesting(self, model, guide, args, kwargs): """ diff --git a/pyro/nn/module.py b/pyro/nn/module.py index cae6a1e8d2..3b335c97b5 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -182,11 +182,13 @@ 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=None) - self._param_ctx.__enter__() + 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 @@ -419,6 +421,8 @@ def named_pyro_params(self, prefix="", recurse=True): yield elem def _pyro_set_supermodule(self, name, context): + if _is_module_local_param_enabled(): + context.param_state.update(self._pyro_context.param_state) self._pyro_name = name self._pyro_context = context for key, value in self._modules.items(): diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index c95f36dd1b..b93b5f2283 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -13,7 +13,14 @@ import pyro.distributions as dist from pyro import poutine from pyro.infer import SVI, Trace_ELBO -from pyro.nn.module import PyroModule, PyroParam, PyroSample, clear, to_pyro_module_ +from pyro.nn.module import ( + PyroModule, + PyroParam, + PyroSample, + clear, + pyro_method, + to_pyro_module_, +) from pyro.optim import Adam from tests.common import assert_equal @@ -132,6 +139,72 @@ def forward(self, data): assert torch.allclose(p, p2) +@pytest.mark.parametrize("local_params", [True, False]) +def test_local_param_vanilla_behavior(local_params): + class Model(PyroModule): + def __init__(self): + super().__init__() + self.nn_param = nn.Parameter(torch.tensor(0.0)) + self.pyro_nn_param = PyroParam( + torch.tensor(1.0), constraint=constraints.positive + ) + + @pyro_method + def init_pyro_param(self): + return pyro.param( + "pyro_param", lambda: torch.tensor(0.5), constraint=constraints.positive + ) + + @pyro_method + def get_pyro_param(self): + return pyro.param("pyro_param") + + @pyro_method + def update_params(self) -> None: + self.get_pyro_param().unconstrained().data += torch.randn(()) + self.nn_param.unconstrained().data += torch.randn(()) + self.pyro_nn_param.unconstrained().data += torch.rand(()) + + def forward(self): + return self.nn_param, self.pyro_nn_param, self.init_pyro_param() + + with pyro.module_local_param_enabled(local_params): + model = Model() + model() # initialize + + pyro_param_0 = model.get_pyro_param() + model.update_params() + pyro_param_1 = model.get_pyro_param() + + assert pyro_param_1 != pyro_param_0 + + model2 = Model() + model2() # initialize + pyro_param_2 = model2.get_pyro_param() + + if local_params: + assert set(pyro.get_param_store().keys()) == set() + assert pyro_param_2 == pyro_param_0 + assert pyro_param_2 != pyro_param_1 + else: + assert set(pyro.get_param_store().keys()) != set() + assert pyro_param_2 != pyro_param_0 + assert pyro_param_2 == pyro_param_1 + + # vanilla param names should be unaffected by module nesting + outer = PyroModule(name="outer") + outer.inner = model + outer.inner() # initialize + pyro_param_3 = outer.inner.get_pyro_param() + + assert pyro_param_3 == pyro_param_1 + + if local_params: + assert "pyro_param" in set(outer._pyro_context.param_state["params"].keys()) + else: + assert "pyro_param" in set(pyro.get_param_store().keys()) + + @pytest.mark.parametrize("local_params", [True, False]) def test_names(local_params): class Model(PyroModule): @@ -620,7 +693,9 @@ def randomize(model): assert_identical(actual, expected) -def test_torch_serialize_attributes(): +@pytest.mark.parametrize("local_params", [True, False]) +def test_torch_serialize_attributes(local_params): + pyro.enable_module_local_param(local_params) module = PyroModule() module.x = PyroParam(torch.tensor(1.234), constraints.positive) module.y = nn.Parameter(torch.randn(3)) From 401e9c61fcc8827645b1206cce41c2b6f7f14b34 Mon Sep 17 00:00:00 2001 From: Eli Date: Mon, 24 Oct 2022 20:54:47 -0400 Subject: [PATCH 13/17] only propagate vanilla pyro.params in the param_state update --- pyro/nn/module.py | 16 +++++++++++++++- tests/nn/test_module.py | 3 ++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/pyro/nn/module.py b/pyro/nn/module.py index 3b335c97b5..d15bc9c63c 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -422,7 +422,21 @@ def named_pyro_params(self, prefix="", recurse=True): def _pyro_set_supermodule(self, name, context): if _is_module_local_param_enabled(): - context.param_state.update(self._pyro_context.param_state) + 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(): diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index b93b5f2283..26b12e7491 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -192,12 +192,13 @@ def forward(self): assert pyro_param_2 == pyro_param_1 # vanilla param names should be unaffected by module nesting - outer = PyroModule(name="outer") + outer = PyroModule[torch.nn.Module]() outer.inner = model outer.inner() # initialize pyro_param_3 = outer.inner.get_pyro_param() assert pyro_param_3 == pyro_param_1 + assert pyro_param_3 != pyro_param_0 if local_params: assert "pyro_param" in set(outer._pyro_context.param_state["params"].keys()) From 446e8a53e7d82f522d06304d149742a45215fb48 Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 1 Nov 2022 16:14:19 -0400 Subject: [PATCH 14/17] settings --- pyro/nn/module.py | 9 +++++++++ pyro/poutine/runtime.py | 4 ---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/pyro/nn/module.py b/pyro/nn/module.py index d15bc9c63c..b25afe087e 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -24,6 +24,15 @@ from pyro.poutine.runtime import _PYRO_PARAM_STORE +@pyro.settings.register("module_local_params", __name__, "MODULE_LOCAL_PARAMS") +def _validate_module_local_params(value: bool) -> None: + assert isinstance(value, bool), "MODULE_LOCAL_PARAMS must be a bool" + + +def _module_local_param_enabled(): + return pyro.settings.get("module_local_params") + + class PyroParam(namedtuple("PyroParam", ("init_value", "constraint", "event_dim"))): """ Declares a Pyro-managed learnable attribute of a :class:`PyroModule`, diff --git a/pyro/poutine/runtime.py b/pyro/poutine/runtime.py index 85c6e19274..f11c89546d 100644 --- a/pyro/poutine/runtime.py +++ b/pyro/poutine/runtime.py @@ -15,10 +15,6 @@ # the global ParamStore _PYRO_PARAM_STORE = ParamStoreDict() -# toggle usage of local param stores in PyroModules -_PYRO_MODULE_LOCAL_PARAM = False - - class _DimAllocator: """ Dimension allocator for internal use by :class:`plate`. From d7c4777f9da798d660165d75ed24f074c465c36d Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 1 Nov 2022 16:25:30 -0400 Subject: [PATCH 15/17] update with settings --- pyro/__init__.py | 2 -- pyro/nn/module.py | 13 ++++++------- pyro/poutine/runtime.py | 1 + pyro/primitives.py | 32 -------------------------------- tests/conftest.py | 2 +- tests/nn/test_module.py | 12 ++++++------ 6 files changed, 14 insertions(+), 48 deletions(-) diff --git a/pyro/__init__.py b/pyro/__init__.py index 858388d72d..bd3b143c58 100644 --- a/pyro/__init__.py +++ b/pyro/__init__.py @@ -9,14 +9,12 @@ 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, diff --git a/pyro/nn/module.py b/pyro/nn/module.py index b25afe087e..0708a6ba02 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -23,13 +23,16 @@ from pyro.ops.provenance import detach_provenance from pyro.poutine.runtime import _PYRO_PARAM_STORE +_MODULE_LOCAL_PARAMS: bool = False -@pyro.settings.register("module_local_params", __name__, "MODULE_LOCAL_PARAMS") + +@pyro.settings.register("module_local_params", __name__, "_MODULE_LOCAL_PARAMS") def _validate_module_local_params(value: bool) -> None: - assert isinstance(value, bool), "MODULE_LOCAL_PARAMS must be a bool" + assert isinstance(value, bool) -def _module_local_param_enabled(): +def _is_module_local_param_enabled() -> bool: + print(pyro.settings.get("module_local_params")) return pyro.settings.get("module_local_params") @@ -178,10 +181,6 @@ 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. diff --git a/pyro/poutine/runtime.py b/pyro/poutine/runtime.py index f11c89546d..59b27c8911 100644 --- a/pyro/poutine/runtime.py +++ b/pyro/poutine/runtime.py @@ -15,6 +15,7 @@ # the global ParamStore _PYRO_PARAM_STORE = ParamStoreDict() + class _DimAllocator: """ Dimension allocator for internal use by :class:`plate`. diff --git a/pyro/primitives.py b/pyro/primitives.py index 0fa34e3ca0..ca67d7b061 100644 --- a/pyro/primitives.py +++ b/pyro/primitives.py @@ -562,35 +562,3 @@ 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 diff --git a/tests/conftest.py b/tests/conftest.py index 5ed178546c..e8b8f2a043 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -34,7 +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) + pyro.settings.set(module_local_params=False) def pytest_addoption(parser): diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 26b12e7491..2db3363172 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -94,7 +94,7 @@ def forward(self, data): with pyro.plate("data"): pyro.sample("obs", dist.Cauchy(loc, log_scale.exp()), obs=data) - with pyro.module_local_param_enabled(local_params): + with pyro.settings.context(module_local_params=local_params): data = torch.randn(5) model = Model() model(data) # initialize @@ -168,7 +168,7 @@ def update_params(self) -> None: def forward(self): return self.nn_param, self.pyro_nn_param, self.init_pyro_param() - with pyro.module_local_param_enabled(local_params): + with pyro.settings.context(module_local_params=local_params): model = Model() model() # initialize @@ -227,7 +227,7 @@ def forward(self): self.p.v self.p.w - with pyro.module_local_param_enabled(local_params): + with pyro.settings.context(module_local_params=local_params): model = Model() # Check named_parameters. @@ -419,7 +419,7 @@ def __init__(self): def forward(self): return [x.clone() for x in [self.x, self.m.weight, self.m.bias, self.p.x]] - with pyro.module_local_param_enabled(local_params): + with pyro.settings.context(module_local_params=local_params): m = Model() state0 = m() @@ -696,7 +696,7 @@ def randomize(model): @pytest.mark.parametrize("local_params", [True, False]) def test_torch_serialize_attributes(local_params): - pyro.enable_module_local_param(local_params) + pyro.settings.set(module_local_params=local_params) module = PyroModule() module.x = PyroParam(torch.tensor(1.234), constraints.positive) module.y = nn.Parameter(torch.randn(3)) @@ -719,7 +719,7 @@ def test_torch_serialize_attributes(local_params): @pytest.mark.parametrize("local_params", [True, False]) def test_torch_serialize_decorators(local_params): - with pyro.module_local_param_enabled(local_params): + with pyro.settings.context(module_local_params=local_params): module = DecoratorModel(3) module() # initialize From 8a42dbf43d4d297e9d67b09ec8a5cf6db37907f9 Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 1 Nov 2022 17:51:17 -0400 Subject: [PATCH 16/17] try to raise an error instead of handling global params --- pyro/nn/module.py | 40 ++++++++++++---------- tests/nn/test_module.py | 73 +++++++---------------------------------- 2 files changed, 33 insertions(+), 80 deletions(-) diff --git a/pyro/nn/module.py b/pyro/nn/module.py index 0708a6ba02..8168461f0c 100644 --- a/pyro/nn/module.py +++ b/pyro/nn/module.py @@ -32,7 +32,6 @@ def _validate_module_local_params(value: bool) -> None: def _is_module_local_param_enabled() -> bool: - print(pyro.settings.get("module_local_params")) return pyro.settings.get("module_local_params") @@ -429,22 +428,8 @@ 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) + if _is_module_local_param_enabled() and pyro.settings.get("validate_poutine"): + self._check_module_local_param_usage() self._pyro_name = name self._pyro_context = context for key, value in self._modules.items(): @@ -460,7 +445,26 @@ def _pyro_get_fullname(self, name): def __call__(self, *args, **kwargs): with self._pyro_context: - return super().__call__(*args, **kwargs) + result = super().__call__(*args, **kwargs) + if ( + pyro.settings.get("validate_poutine") + and not self._pyro_context.active + and _is_module_local_param_enabled() + ): + self._check_module_local_param_usage() + return result + + def _check_module_local_param_usage(self) -> None: + self_nn_params = set(id(p) for p in self.parameters()) + self_pyro_params = set( + id(p if not hasattr(p, "unconstrained") else p.unconstrained()) + for p in self._pyro_context.param_state["params"].values() + ) + if not self_pyro_params <= self_nn_params: + raise NotImplementedError( + "Support for global pyro.param statements in PyroModules " + "with local param mode enabled is not yet implemented." + ) def __getattr__(self, name): # PyroParams trigger pyro.param statements. diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 2db3363172..31941f1bd5 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -13,14 +13,7 @@ import pyro.distributions as dist from pyro import poutine from pyro.infer import SVI, Trace_ELBO -from pyro.nn.module import ( - PyroModule, - PyroParam, - PyroSample, - clear, - pyro_method, - to_pyro_module_, -) +from pyro.nn.module import PyroModule, PyroParam, PyroSample, clear, to_pyro_module_ from pyro.optim import Adam from tests.common import assert_equal @@ -140,70 +133,26 @@ def forward(self, data): @pytest.mark.parametrize("local_params", [True, False]) -def test_local_param_vanilla_behavior(local_params): +def test_local_param_global_behavior_fails(local_params): class Model(PyroModule): def __init__(self): super().__init__() - self.nn_param = nn.Parameter(torch.tensor(0.0)) - self.pyro_nn_param = PyroParam( - torch.tensor(1.0), constraint=constraints.positive - ) - - @pyro_method - def init_pyro_param(self): - return pyro.param( - "pyro_param", lambda: torch.tensor(0.5), constraint=constraints.positive - ) - - @pyro_method - def get_pyro_param(self): - return pyro.param("pyro_param") - - @pyro_method - def update_params(self) -> None: - self.get_pyro_param().unconstrained().data += torch.randn(()) - self.nn_param.unconstrained().data += torch.randn(()) - self.pyro_nn_param.unconstrained().data += torch.rand(()) + self.global_nn_param = nn.Parameter(torch.zeros(2)) def forward(self): - return self.nn_param, self.pyro_nn_param, self.init_pyro_param() + global_param = pyro.param("_global_param", lambda: torch.randn(2)) + global_nn_param = pyro.param("global_nn_param", self.global_nn_param) + return global_param, global_nn_param with pyro.settings.context(module_local_params=local_params): model = Model() - model() # initialize - - pyro_param_0 = model.get_pyro_param() - model.update_params() - pyro_param_1 = model.get_pyro_param() - - assert pyro_param_1 != pyro_param_0 - - model2 = Model() - model2() # initialize - pyro_param_2 = model2.get_pyro_param() - if local_params: - assert set(pyro.get_param_store().keys()) == set() - assert pyro_param_2 == pyro_param_0 - assert pyro_param_2 != pyro_param_1 - else: - assert set(pyro.get_param_store().keys()) != set() - assert pyro_param_2 != pyro_param_0 - assert pyro_param_2 == pyro_param_1 - - # vanilla param names should be unaffected by module nesting - outer = PyroModule[torch.nn.Module]() - outer.inner = model - outer.inner() # initialize - pyro_param_3 = outer.inner.get_pyro_param() - - assert pyro_param_3 == pyro_param_1 - assert pyro_param_3 != pyro_param_0 - - if local_params: - assert "pyro_param" in set(outer._pyro_context.param_state["params"].keys()) + assert pyro.settings.get("module_local_params") + with pytest.raises(NotImplementedError): + model() else: - assert "pyro_param" in set(pyro.get_param_store().keys()) + assert not pyro.settings.get("module_local_params") + model() @pytest.mark.parametrize("local_params", [True, False]) From 41a157c1fd2594b9067ce29e43485f2dc35b5f9a Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 8 Nov 2022 22:25:39 -0500 Subject: [PATCH 17/17] avoid use of settings.set in tests and remove change in global conftest --- tests/conftest.py | 1 - tests/nn/test_module.py | 34 +++++++++++++++++----------------- 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index e8b8f2a043..4103356fe3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -34,7 +34,6 @@ def pytest_runtest_setup(item): if test_initialize_marker: rng_seed = test_initialize_marker.kwargs["rng_seed"] pyro.set_rng_seed(rng_seed) - pyro.settings.set(module_local_params=False) def pytest_addoption(parser): diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py index 31941f1bd5..90b8af0eec 100644 --- a/tests/nn/test_module.py +++ b/tests/nn/test_module.py @@ -645,25 +645,25 @@ def randomize(model): @pytest.mark.parametrize("local_params", [True, False]) def test_torch_serialize_attributes(local_params): - pyro.settings.set(module_local_params=local_params) - module = PyroModule() - module.x = PyroParam(torch.tensor(1.234), constraints.positive) - module.y = nn.Parameter(torch.randn(3)) - assert isinstance(module.x, torch.Tensor) + with pyro.settings.context(module_local_params=local_params): + module = PyroModule() + module.x = PyroParam(torch.tensor(1.234), constraints.positive) + module.y = nn.Parameter(torch.randn(3)) + assert isinstance(module.x, torch.Tensor) - # Work around https://github.com/pytorch/pytorch/issues/27972 - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=UserWarning) - f = io.BytesIO() - torch.save(module, f) - pyro.clear_param_store() - f.seek(0) - actual = torch.load(f) + # Work around https://github.com/pytorch/pytorch/issues/27972 + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=UserWarning) + f = io.BytesIO() + torch.save(module, f) + pyro.clear_param_store() + f.seek(0) + actual = torch.load(f) - assert_equal(actual.x, module.x) - actual_names = {name for name, _ in actual.named_parameters()} - expected_names = {name for name, _ in module.named_parameters()} - assert actual_names == expected_names + assert_equal(actual.x, module.x) + actual_names = {name for name, _ in actual.named_parameters()} + expected_names = {name for name, _ in module.named_parameters()} + assert actual_names == expected_names @pytest.mark.parametrize("local_params", [True, False])