Get mcmc sampling to work #9

sharanry · 2018-06-13T18:37:08Z

Unobserved variables accessible
Sampling works

TODO:
Write Tests

sharanry · 2018-06-13T18:53:55Z

@ferrine
What do you suggest model.target_log_prob_fn() give? The logp function of only the result of model.f or all the intermediate RVs too?

If we give it only the final logp function then we wont be able to sample traces of intermediate RVs.
One more problem is tfp.mcmc.sample_chain() might work for only single RV(in current implementation, final) logp function .

And should the model.unobserved contain only final RV, i.e, the result of f()?

ferrine

Review

Graph namespace

TF uses graph.as_default() to create graph. So if you repeatedly call model.unobserved it will spoil the namespace totally. The below snippet can replicate the problem

import tensorflow as tf
graph = tf.get_default_graph()
sess = tf.InteractiveSession(graph=graph)
def model():
    return tf.ones([1])
model()
model()
model()
graph.as_graph_def()

The output contains a lot of versions of tf.ones(). One way to solve the problem it to put all internal things into an auxiliary namespace.

ferrine · 2018-06-15T10:47:38Z

pymc4/inference/sampling/sample.py

-    for name, shape in model.unobserved.iteritems():
-        initial_state.append(.5 * tf.ones(shape, name="init_{}".format(name)))
+    for name in model.unobserved:
+        initial_state.append(.5 * tf.ones(model.unobserved[name].shape, name="init_{}".format(name)))


This will create a lot of problems with namespace

I do not know the way to go but to avoid repeated calls of self._f

unobserved = {} for i in self.variables: if self.variables[i] not in self.observed.values(): unobserved[i] = self.variables[i] unobserved = collections.OrderedDict(unobserved) return unobserved

I could do this to avoid unobserved calling f() multiple times?

In [2]: graph = tf.Graph() ...: ...: with graph.as_default(): ...: ed.Normal(0., 1.) ...: print(graph.as_graph_def()) ...:

Outputs:

node { name: "Normal/loc/input" op: "Const" attr { key: "dtype" value { type: DT_FLOAT } } attr { key: "value" value { tensor { dtype: DT_FLOAT tensor_shape { } float_val: 0.0 } } } } node { name: "Normal/loc" op: "Identity" input: "Normal/loc/input" attr { key: "T" value { type: DT_FLOAT } } } node { name: "Normal/scale/input" op: "Const" attr { key: "dtype" value { type: DT_FLOAT } } attr { key: "value" value { tensor { dtype: DT_FLOAT tensor_shape { } float_val: 1.0 } } } } node { name: "Normal/scale" op: "Identity" input: "Normal/scale/input" attr { key: "T" value { type: DT_FLOAT } } } node { name: "Normal_1/sample/sample_shape" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { } } } } } } node { name: "Normal_1/sample/Normal/batch_shape_tensor/batch_shape" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { } } } } } } node { name: "Normal_1/sample/concat/values_0" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 1 } } } } node { name: "Normal_1/sample/concat/axis" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { } int_val: 0 } } } } node { name: "Normal_1/sample/concat" op: "ConcatV2" input: "Normal_1/sample/concat/values_0" input: "Normal_1/sample/Normal/batch_shape_tensor/batch_shape" input: "Normal_1/sample/concat/axis" attr { key: "N" value { i: 2 } } attr { key: "T" value { type: DT_INT32 } } attr { key: "Tidx" value { type: DT_INT32 } } } node { name: "Normal_1/sample/random_normal/mean" op: "Const" attr { key: "dtype" value { type: DT_FLOAT } } attr { key: "value" value { tensor { dtype: DT_FLOAT tensor_shape { } float_val: 0.0 } } } } node { name: "Normal_1/sample/random_normal/stddev" op: "Const" attr { key: "dtype" value { type: DT_FLOAT } } attr { key: "value" value { tensor { dtype: DT_FLOAT tensor_shape { } float_val: 1.0 } } } } node { name: "Normal_1/sample/random_normal/RandomStandardNormal" op: "RandomStandardNormal" input: "Normal_1/sample/concat" attr { key: "T" value { type: DT_INT32 } } attr { key: "dtype" value { type: DT_FLOAT } } attr { key: "seed" value { i: 0 } } attr { key: "seed2" value { i: 0 } } } node { name: "Normal_1/sample/random_normal/mul" op: "Mul" input: "Normal_1/sample/random_normal/RandomStandardNormal" input: "Normal_1/sample/random_normal/stddev" attr { key: "T" value { type: DT_FLOAT } } } node { name: "Normal_1/sample/random_normal" op: "Add" input: "Normal_1/sample/random_normal/mul" input: "Normal_1/sample/random_normal/mean" attr { key: "T" value { type: DT_FLOAT } } } node { name: "Normal_1/sample/mul" op: "Mul" input: "Normal_1/sample/random_normal" input: "Normal/scale" attr { key: "T" value { type: DT_FLOAT } } } node { name: "Normal_1/sample/add" op: "Add" input: "Normal_1/sample/mul" input: "Normal/loc" attr { key: "T" value { type: DT_FLOAT } } } node { name: "Normal_1/sample/Shape" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 1 } } } } node { name: "Normal_1/sample/strided_slice/stack" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 1 } } } } node { name: "Normal_1/sample/strided_slice/stack_1" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 0 } } } } node { name: "Normal_1/sample/strided_slice/stack_2" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 1 } } } } node { name: "Normal_1/sample/strided_slice" op: "StridedSlice" input: "Normal_1/sample/Shape" input: "Normal_1/sample/strided_slice/stack" input: "Normal_1/sample/strided_slice/stack_1" input: "Normal_1/sample/strided_slice/stack_2" attr { key: "Index" value { type: DT_INT32 } } attr { key: "T" value { type: DT_INT32 } } attr { key: "begin_mask" value { i: 0 } } attr { key: "ellipsis_mask" value { i: 0 } } attr { key: "end_mask" value { i: 1 } } attr { key: "new_axis_mask" value { i: 0 } } attr { key: "shrink_axis_mask" value { i: 0 } } } node { name: "Normal_1/sample/concat_1/axis" op: "Const" attr { key: "dtype" value { type: DT_INT32 } } attr { key: "value" value { tensor { dtype: DT_INT32 tensor_shape { } int_val: 0 } } } } node { name: "Normal_1/sample/concat_1" op: "ConcatV2" input: "Normal_1/sample/sample_shape" input: "Normal_1/sample/strided_slice" input: "Normal_1/sample/concat_1/axis" attr { key: "N" value { i: 2 } } attr { key: "T" value { type: DT_INT32 } } attr { key: "Tidx" value { type: DT_INT32 } } } node { name: "Normal_1/sample/Reshape" op: "Reshape" input: "Normal_1/sample/add" input: "Normal_1/sample/concat_1" attr { key: "T" value { type: DT_FLOAT } } attr { key: "Tshape" value { type: DT_INT32 } } } versions { producer: 26 }

For just one edward random variable there is so much change in the graph.

That may be okay, if graph modifications are not frequently called (or it is hard to do) by user.

ferrine

There is a special interceptor for this purpose

ferrine · 2018-06-21T08:05:30Z

pymc4/model/base.py

@@ -68,14 +68,41 @@ def get_mode(state, rv, *args, **kwargs):
            returns = self.session.run(list(values_collector.result.values()))
        return dict(zip(values_collector.result.keys(), returns))

-    def target_log_prob_fn(self, *args, **kwargs):
+    def log_prob_fn(self, x,  *args, **kwargs):


what is x for here?

Not necessary, removing it.

ferrine · 2018-06-21T08:05:34Z

pymc4/model/base.py

+        def log_joint_fn(*args, **kwargs):
+            states = dict(zip(self.unobserved.keys(), args))
+            states.update(self.observed)
+            log_probs = []


https://github.com/pymc-devs/pymc4/blob/functional/pymc4/util/interceptors.py#L110

collect_log_prob = CollectLogProb(states) with ed.interception(collect_log_prob): self._f(self._cfg) return collect_log_prob.result

changing it, was facing problems with states before. Now working.

@ferrine interceptors.CollectLogProb only works with model like

@model.define def process(cfg=None): mu = ed.Normal(0., 1., name="mu") obs = ed.Normal(0., 1., name="obs") return obs

and not model like

@model.define def process(cfg=None): mu = ed.Normal(0., 1., name="mu") obs = ed.Normal(mu, 1., name="obs") return obs

hmm, what's happening?

ferrine · 2018-06-21T08:09:39Z

pymc4/model/base.py

@@ -67,6 +68,41 @@ def get_mode(state, rv, *args, **kwargs):
            returns = self.session.run(list(values_collector.result.values()))
        return dict(zip(values_collector.result.keys(), returns))

+    def log_prob_fn(self, x,  *args, **kwargs):


hmm, I'm not sure this will work. ancestors of RV depend on the RV, here you do not replace RV with value=kwargs.get(i)

ColCarroll

Just a few styling nitpicks around dictionary iteration!

ColCarroll · 2018-06-24T14:50:30Z

pymc4/inference/sampling/sample.py

@@ -10,7 +10,8 @@ def sample(model,
           num_leapfrog_steps=3,
           numpy=True):
    initial_state = []
-    for name, shape in model.unobserved.iteritems():
+    for name in model.unobserved.keys():


use for name, (_, shape, _) in model.items(): to indicate that dist and rv are not used in the loop

ColCarroll · 2018-06-24T14:53:46Z

pymc4/model/base.py

+    @property
+    def unobserved(self):
+        unobserved = {}
+        for i in self.variables:


for name, variable in self.variables.items(): if variable not in self.observed.values(): unobserved[name] = variable

ColCarroll · 2018-06-24T15:00:05Z

pymc4/model/base.py

@@ -83,6 +100,16 @@ def graph(self):
    def observed(self):
        return self._observed

+    @property
+    def unobserved(self):
+        unobserved = {}


this can be initialized as an OrderedDict: currently in Python < 3.6, the return value will not be ordered, since you built an (unordered) dictionary, then turned it into an OrderedDict. We're targeting 3.6 and higher though, in which case you do not need the OrderedDict at all, since all dicts now maintain insertion order.

That's a long way to say: I would just make a plain dictionary, but if you use OrderedDict, it needs to be initialized as such.

ferrine · 2018-06-26T08:10:16Z

pymc4/util/interceptors.py

@@ -9,7 +9,7 @@
    'CollectLogProb'
 ]

-VariableDescription = collections.namedtuple('VariableDescription', 'Dist,shape')
+VariableDescription = collections.namedtuple('VariableDescription', 'Dist,shape,rv')


I'm still worried about this solution. Variable description is supposed to be collected far before sampling (or what about changing this?). So when you first collect VariableInfo for the first time and get these RVs, you save temporary nodes. When you collect LogProb you again run the model and variables involves there are totally different from those that are stored in VariableDescription. That's why I did not store them there.

Agree with @ferrine - The RVs are not initialized until we configure the model (following the idea in the API discussion doc, we create the RVs when we call model.configure(...) or model.sample(...)). This means that we record the Distribution and the relationship between RVs, but the actually RVs are only initialized when we actually using them (ie, in the evaluation of logp).

The problem i am facing if RVs are not stored in the VariableDescription is that it doesn't store the specifics of any distribution (like loc or scale) even if they are mentioned in the model definition. So we will have to collect all this already provided info somehow.

@model.define def process(): mu = ed.Normal(loc=0., scale=10., name="mu") # here we lose the info that it has loc 0 and scale 10 without RV.

We can try defining a new Interceptor which does this for us for each RV.
We can then overwrite(replace existing and add new) the collected data every-time we call configure.

OK, let's make it later

sharanry · 2018-07-03T13:37:05Z

@ferrine Could I merge this PR?

ferrine · 2018-07-04T07:14:50Z

pymc4/model/base.py

+            states.update(self.observed)
+            log_probs = []
+
+            def interceptor(f, *args, **kwargs):


Don't we want use class based interceptor for consistency?

ferrine · 2018-07-04T07:16:24Z

pymc4/util/interceptors.py

@@ -9,7 +9,7 @@
    'CollectLogProb'
 ]

-VariableDescription = collections.namedtuple('VariableDescription', 'Dist,shape')
+VariableDescription = collections.namedtuple('VariableDescription', 'Dist,shape,rv')


OK, let's make it later

ferrine · 2018-07-04T07:17:42Z

tests/test_model.py

+    assert len(model.observed) == 1
+    assert not model.unobserved
+
+    model.reset()


We decided to meke a copy of model each time state changes

this one is not critical though, refactoring interceptor usage is what is really needed to finish this PR (#9 (diff))

sharanry · 2018-07-06T06:16:39Z

@ferrine I have changed it to a class based interceptor

ferrine

I can't find sampling test, I think we need one.

ferrine · 2018-07-06T07:15:12Z

And test point is better get via model.test_point()

sharanry · 2018-07-06T09:20:36Z

@ferrine Currently model.test_point() is how you get the test point.
I am not sure I understand what you are saying.

ferrine · 2018-07-06T10:00:09Z

pymc4/inference/sampling/sample.py

           num_results=5000,
           num_burnin_steps=3000,
           step_size=.4,
           num_leapfrog_steps=3,
           numpy=True):
    initial_state = []
-    for name, shape in model.unobserved.iteritems():
+    for name, (_, shape, _) in model.unobserved.items():


for name, point in model.test_point(mode=mode): initial_state.append(point)

may be done in next PR

ColCarroll · 2018-07-06T11:27:18Z

Congrats @sharanry , and thanks for the thorough review @ferrine !

sharanry requested review from ferrine and removed request for ferrine June 13, 2018 18:37

sharanry changed the title ~~Get mcmc sampling to work~~ [WIP] Get mcmc sampling to work Jun 13, 2018

sharanry self-assigned this Jun 13, 2018

ferrine reviewed Jun 15, 2018

View reviewed changes

Add model.target_log_prob_fn()

a703c21

sharanry force-pushed the unobserved branch from 00651f6 to a703c21 Compare June 18, 2018 11:48

add target_log_prob_fn which works with the tff mcmc sampler

66f95ca

ferrine reviewed Jun 21, 2018

View reviewed changes

sharanry added 2 commits June 23, 2018 12:16

modify target_log_prob_fn

8aaa0ff

remove log_prob_fn

cded7c5

ColCarroll reviewed Jun 24, 2018

View reviewed changes

ferrine reviewed Jun 26, 2018

View reviewed changes

sharanry added 9 commits June 26, 2018 19:14

minor fixes

c27e97f

add some tests

a8e3dae

add tests and fix model.configure()

4f5382f

fix lint errors

e06d946

add tests

ca9f334

remove some tests

a7cef9b

fix pycodestyle errors

bd381b1

remove test_interceptors

89edc5c

add pycodestyle to dev requirements

9f46878

sharanry changed the title ~~[WIP] Get mcmc sampling to work~~ Get mcmc sampling to work Jul 1, 2018

ferrine suggested changes Jul 4, 2018

View reviewed changes

make target_log_prob_fn to use class based interceptor

f223e4e

fix pylint error

6717b3e

ferrine reviewed Jul 6, 2018

View reviewed changes

sharanry added 2 commits July 6, 2018 14:42

add test for sampling

007cd06

solve pycodestyle errors

d896372

ferrine reviewed Jul 6, 2018

View reviewed changes

ferrine approved these changes Jul 6, 2018

View reviewed changes

ColCarroll merged commit 4c8d0d5 into pymc-devs:functional Jul 6, 2018

Get mcmc sampling to work #9

Get mcmc sampling to work #9

Conversation

sharanry commented Jun 13, 2018 • edited Loading

sharanry commented Jun 13, 2018 • edited Loading

ferrine left a comment

Choose a reason for hiding this comment

Review

Graph namespace

Choose a reason for hiding this comment

ferrine Jun 15, 2018 • edited Loading

Choose a reason for hiding this comment

sharanry Jun 16, 2018 • edited Loading

Choose a reason for hiding this comment

sharanry Jun 16, 2018 • edited by ferrine Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ferrine left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ColCarroll left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sharanry commented Jul 3, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sharanry commented Jul 6, 2018

ferrine left a comment

Choose a reason for hiding this comment

ferrine commented Jul 6, 2018

sharanry commented Jul 6, 2018

Choose a reason for hiding this comment

ferrine Jul 6, 2018 • edited Loading

Choose a reason for hiding this comment

ColCarroll commented Jul 6, 2018

sharanry commented Jun 13, 2018 •

edited

Loading

sharanry commented Jun 13, 2018 •

edited

Loading

ferrine Jun 15, 2018 •

edited

Loading

sharanry Jun 16, 2018 •

edited

Loading

sharanry Jun 16, 2018 •

edited by ferrine

Loading

ferrine Jul 6, 2018 •

edited

Loading