def test_only_withs():
def model1():
with scope(prefix="a"):
with scope(prefix="b"):
pyro.sample("x", dist.Bernoulli(0.5))
tr1 = poutine.trace(name_count(model1)).get_trace()
assert "a/b/x" in tr1.nodes
tr2 = poutine.trace(name_count(scope(prefix="model1")(model1))).get_trace()
assert "model1/a/b/x" in tr2.nodes
def test_recur_multi():
@scope(inner=True)
def model1(r=True):
model2()
with scope(prefix="inter"):
model2()
if r:
model1(r=False)
model2()
@scope(inner=True)
def model2():
return pyro.sample("y", dist.Normal(0.0, 1.0))
true_samples = ["model1/model2/y",
"model1/inter/model2/y",
"model1/inter/model1/model2/y",
"model1/inter/model1/inter/model2/y",
"model1/inter/model1/model2/y__1",
"model1/model2/y__1"]
tr = poutine.trace(name_count(model1)).get_trace()
samples = [name for name, node in tr.nodes.items()
if node["type"] == "sample"]
logger.debug(samples)
assert true_samples == samples
def infer_and_run(
model,
num_samples=5000,
num_iterations=2000,
debug=False,
learning_rate=0.01,
early_stopping_patience=200,
) -> pd.DataFrame:
"""
debug - whether to output debug information
num_iterations - Number of optimizer iterations
learning_rate - Optimizer learning rate
early_stopping_patience - Stop training if loss hasn't improved for this many iterations
"""
def to_numpy(d):
return {k: v.detach().numpy() for k, v in d.items()}
def debug_output(guide):
quantiles = to_numpy(guide.quantiles([0.05, 0.5, 0.95]))
for k, v in quantiles.items():
print(f"{k}: {v[1]:.4f} [{v[0]:.4f}, {v[2]:.4f}]")
model = name_count(model)
# Automatically chooses a normal distribution for each variable
guide = pyro.infer.autoguide.AutoNormal(
model, init_loc_fn=pyro.infer.autoguide.init_to_median)
pyro.clear_param_store()
if debug:
guide(training=True) # Needed to initialize the guide before output
debug_output(guide)
print()
adam = pyro.optim.Adam({"lr": learning_rate})
svi = SVI(model, guide, adam, loss=Trace_ELBO())
best_loss = None
last_improvement = None
for j in range(num_iterations):
# calculate the loss and take a gradient step
loss = svi.step(training=True)
if best_loss is None or best_loss > loss:
best_loss = loss
last_improvement = j
if j % 100 == 0:
if debug:
print("[iteration %04d]" % (j + 1))
print(f"loss: {loss:.4f}")
debug_output(guide)
print()
if j > (last_improvement + early_stopping_patience):
print("Stopping Early")
break
print(f"Final loss: {loss:.4f}")
predictive = Predictive(model, guide=guide, num_samples=num_samples)
raw_samples = predictive(training=False)
return pd.DataFrame(to_numpy(raw_samples))
def test_nested_traces():
@scope
def f1():
return pyro.sample("x", dist.Bernoulli(0.5))
@scope
def f2():
f1()
f1()
f1()
return pyro.sample("y", dist.Bernoulli(0.5))
expected_names = ["f2/f1/x", "f2/f1__1/x", "f2/f1__2/x", "f2/y"]
tr2 = poutine.trace(name_count(name_count(f2))).get_trace()
actual_names = [name for name, node in tr2.nodes.items()
if node['type'] == "sample"]
assert expected_names == actual_names
def sample_morphism(self, obj, probs, temperature, min_depth=2, infer={}):
with name_count():
if obj in self._graph.nodes:
return self.path_between(Ty(), obj, probs, temperature,
min_depth, infer)
entries = unification.unfold_arrow(obj)
src, dest = unification.fold_product(entries[:-1]), entries[-1]
return self.path_between(src, dest, probs, temperature, min_depth,
infer)
def run(model, num_samples=5000, ignore_unnamed=True) -> pd.DataFrame:
"""
1. Run model forward, record samples for variables
2. Return dataframe with one row for each execution
"""
model = name_count(model)
samples: List[Dict[str, float]] = []
for _ in tqdm(range(num_samples)):
sample: Dict[str, float] = {}
trace = pyro.poutine.trace(model).get_trace()
for name in trace.nodes.keys():
if trace.nodes[name]["type"] == "sample":
if not ignore_unnamed or not name.startswith("_var"):
sample[name] = trace.nodes[name]["value"].item()
samples.append(sample)
return pd.DataFrame(samples) # type: ignore
def test_simple_recur():
@scope
def geometric(p):
x = pyro.sample("x", dist.Bernoulli(p))
if x.item() == 1.0:
# model1()
return x + geometric(p)
else:
return x
prev_name = "x"
for name, node in poutine.trace(name_count(geometric)).get_trace(0.9).nodes.items():
if node["type"] == "sample":
logger.debug(name)
assert name == "geometric/" + prev_name
prev_name = "geometric/" + prev_name
def set_mode(self, mode):
"""
Sets ``mode`` of this object to be able to use its parameters in
stochastic functions. If ``mode="model"``, a parameter will get its
value from its prior. If ``mode="guide"``, the value will be drawn from
its guide.
.. note:: This method automatically sets ``mode`` for submodules which
belong to :class:`Parameterized` class.
:param str mode: Either "model" or "guide".
"""
with autoname.name_count():
for module in self.modules():
if isinstance(module, Parameterized):
module.mode = mode
def test_mutual_recur():
@scope
def model1(n):
pyro.sample("a", dist.Bernoulli(0.5))
if n <= 0:
return
else:
return model2(n-1)
@scope
def model2(n):
pyro.sample("b", dist.Bernoulli(0.5))
if n <= 0:
return
else:
model1(n)
names = set(["_INPUT", "_RETURN",
"model2/b", "model2/model1/a", "model2/model1/model2/b"])
tr_names = set([name for name in poutine.trace(name_count(model2)).get_trace(1)])
assert names == tr_names
