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