Reuse allocated arrays in FFBS sampler #81

xjing76 · 2021-03-16T22:32:55Z

Attempts to improve FFB by storing the likelihood series within sampler.

I feel like this is probably not as simple as this but hopefully is this the right idea?

If log_lik_fn is subject to change across each step then this would fail.

closes #76

brandonwillard · 2021-03-18T00:59:41Z

pymc3_hmm/step_methods.py

-        # log_lik_t = log_lik_fn(state_seqs)
-        log_lik_t = np.stack([log_lik_fn(np.broadcast_to(m, N)) for m in range(M)])
+        if self.log_lik_t is not None:
+            log_lik_t = self.log_lik_t


This is the right idea, but we still need to recompute the log-likelihood on every call to this function, because the parameters we're sampling will change on each iteration, which—in turn—makes log_lik_fn produce new values.

The question is: how do we recompute the log_lik_fn steps below in-place (or as close to in-place as possible)?

In other words, log_lik_fn is allocating and filling a large length N array, but, now, we want it to fill an existing pre-allocated array with the values it computes.

The answer most likely involves creating our own log_lik_fn function, instead of using the one provided by PyMC3. Ultimately, log_lik_fn is just a compiled Theano/Aesara function, so we need to in-place updating at that level first; however, that might be pretty straightforward, because I believe we can create a shared variable that log_lik_fn updates in-place.

For that matter, we could turn the entire computation (i.e. the stack and range) into a Theano/Aesara graph, use a shared variable and have it updated in-place, then call that function instead of log_lik_fn and pass ffbs_astep the shared variable's underlying value.

xjing76 · 2021-03-19T14:41:46Z

Looking through the Theano graph of how we computed the log likelihood plot.
I think there are a few places we are reallocating the space for sub distributions.

SwitchingProcess I am not sure if we are actually allocating space for logp or we are still broadcasting. However, this seems like a place where we can improved by reusing a shared variables for logp within SwitchingProcess

xjing76 · 2021-03-22T20:34:52Z

Sum{axis=[0], acc_dtype=float64} [id A] ''   
 |InplaceDimShuffle{x,0} [id B] ''   
   |AdvancedIncSubtensor{inplace=False,  set_instead_of_inc=True} [id C] ''   
     |AdvancedIncSubtensor{inplace=False,  set_instead_of_inc=True} [id D] ''   
     | |Alloc [id E] ''   
     | | |TensorConstant{-inf} [id F]
     | | |TensorConstant{30} [id G]
     | |Elemwise{switch,no_inplace} [id H] ''   
     | | |Elemwise{mul,no_inplace} [id I] ''   
     | | | |InplaceDimShuffle{x} [id J] ''   
     | | | | |TensorConstant{1} [id K]
     | | | |Elemwise{mul,no_inplace} [id L] ''   
     | | |   |InplaceDimShuffle{x} [id M] ''   
     | | |   | |TensorConstant{1} [id N]
     | | |   |Elemwise{eq,no_inplace} [id O] ''   
     | | |     |AdvancedSubtensor [id P] ''   
     | | |     | |TensorConstant{[  0   0  ..0   0   0]} [id Q]
     | | |     | |Elemwise{eq,no_inplace} [id R] ''   
     | | |     |   |Elemwise{Cast{int32}} [id S] ''   
     | | |     |   | |S_t ~ DiscreteMarkovChain [id T]
     | | |     |   |InplaceDimShuffle{x} [id U] ''   
     | | |     |     |TensorConstant{0} [id V]
     | | |     |AdvancedSubtensor [id W] ''   
     | | |       |Elemwise{mul,no_inplace} [id X] ''   
     | | |       | |InplaceDimShuffle{x} [id Y] ''   
     | | |       | | |TensorConstant{0} [id Z]
     | | |       | |Alloc [id BA] ''   
     | | |       |   |TensorConstant{1.0} [id BB]
     | | |       |   |Subtensor{int64} [id BC] ''   
     | | |       |     |Shape [id BD] ''   
     | | |       |     | |TensorConstant{[  0   0  ..0   0   0]} [id Q]
     | | |       |     |Constant{0} [id BE]
     | | |       |Elemwise{eq,no_inplace} [id R] ''   
     | | |InplaceDimShuffle{x} [id BF] ''   
     | | | |TensorConstant{0} [id BG]
     | | |InplaceDimShuffle{x} [id BH] ''   
     | |   |TensorConstant{-inf} [id BI]
     | |Elemwise{eq,no_inplace} [id R] ''   
     |Elemwise{switch,no_inplace} [id BJ] ''   
     | |Elemwise{mul,no_inplace} [id BK] ''   
     | | |Elemwise{eq,no_inplace} [id BL] ''   
     | | | |AdvancedSubtensor [id BM] ''   
     | | | | |Elemwise{mul,no_inplace} [id BN] ''   
     | | | | | |InplaceDimShuffle{x} [id BO] ''   
     | | | | | | |TensorConstant{9.0} [id BP]
     | | | | | |Alloc [id BQ] ''   
     | | | | |   |TensorConstant{1.0} [id BR]
     | | | | |   |Subtensor{int64} [id BS] ''   
     | | | | |     |Shape [id BT] ''   
     | | | | |     | |TensorConstant{[  0   0  ..0   0   0]} [id Q]
     | | | | |     |Constant{0} [id BU]
     | | | | |Elemwise{eq,no_inplace} [id BV] ''   
     | | | |   |Elemwise{Cast{int32}} [id S] ''   
     | | | |   |InplaceDimShuffle{x} [id BW] ''   
     | | | |     |TensorConstant{1} [id BX]
     | | | |InplaceDimShuffle{x} [id BY] ''   
     | | |   |TensorConstant{0} [id BZ]
     | | |Elemwise{eq,no_inplace} [id CA] ''   
     | |   |AdvancedSubtensor [id CB] ''   
     | |   | |TensorConstant{[  0   0  ..0   0   0]} [id Q]
     | |   | |Elemwise{eq,no_inplace} [id BV] ''   
     | |   |InplaceDimShuffle{x} [id CC] ''   
     | |     |TensorConstant{0} [id CD]
     | |InplaceDimShuffle{x} [id CE] ''   
     | | |TensorConstant{0} [id CF]
     | |Elemwise{switch,no_inplace} [id CG] ''   
     |   |Elemwise{mul,no_inplace} [id CH] ''   
     |   | |Elemwise{mul,no_inplace} [id CI] ''   
     |   | | |InplaceDimShuffle{x} [id CJ] ''   
     |   | | | |TensorConstant{1} [id CK]
     |   | | |Elemwise{mul,no_inplace} [id CL] ''   
     |   | |   |InplaceDimShuffle{x} [id CM] ''   
     |   | |   | |TensorConstant{1} [id CN]
     |   | |   |Elemwise{ge,no_inplace} [id CO] ''   
     |   | |     |AdvancedSubtensor [id BM] ''   
     |   | |     |InplaceDimShuffle{x} [id CP] ''   
     |   | |       |TensorConstant{0} [id CQ]
     |   | |Elemwise{mul,no_inplace} [id CR] ''   
     |   |   |InplaceDimShuffle{x} [id CS] ''   
     |   |   | |TensorConstant{1} [id CT]
     |   |   |Elemwise{ge,no_inplace} [id CU] ''   
     |   |     |AdvancedSubtensor [id CB] ''   
     |   |     |InplaceDimShuffle{x} [id CV] ''   
     |   |       |TensorConstant{0} [id CW]
     |   |Elemwise{sub,no_inplace} [id CX] ''   
     |   | |Elemwise{sub,no_inplace} [id CY] ''   
     |   | | |Elemwise{switch,no_inplace} [id CZ] ''   
     |   | | | |Elemwise{eq,no_inplace} [id DA] ''   
     |   | | | | |AdvancedSubtensor [id BM] ''   
     |   | | | | |InplaceDimShuffle{x} [id DB] ''   
     |   | | | |   |TensorConstant{0} [id DC]
     |   | | | |Elemwise{switch,no_inplace} [id DD] ''   
     |   | | | | |Elemwise{eq,no_inplace} [id DE] ''   
     |   | | | | | |AdvancedSubtensor [id CB] ''   
     |   | | | | | |InplaceDimShuffle{x} [id DF] ''   
     |   | | | | |   |TensorConstant{0} [id DG]
     |   | | | | |InplaceDimShuffle{x} [id DH] ''   
     |   | | | | | |TensorConstant{0.0} [id DI]
     |   | | | | |InplaceDimShuffle{x} [id DJ] ''   
     |   | | | |   |TensorConstant{-inf} [id DK]
     |   | | | |Elemwise{mul,no_inplace} [id DL] ''   
     |   | | |   |AdvancedSubtensor [id CB] ''   
     |   | | |   |Elemwise{log,no_inplace} [id DM] ''   
     |   | | |     |AdvancedSubtensor [id BM] ''   
     |   | | |gammaln [id DN] ''   
     |   | |   |Elemwise{add,no_inplace} [id DO] ''   
     |   | |     |AdvancedSubtensor [id CB] ''   
     |   | |     |InplaceDimShuffle{x} [id DP] ''   
     |   | |       |TensorConstant{1} [id DQ]
     |   | |AdvancedSubtensor [id BM] ''   
     |   |InplaceDimShuffle{x} [id DR] ''   
     |     |TensorConstant{-inf} [id DS]
     |Elemwise{eq,no_inplace} [id BV] ''

xjing76 · 2021-03-23T21:09:52Z

I am not sure if I understand right. But I think we would like to duplicate the computation graph for each state.
My attempt here However, it seems like when we are cloning the node. We are probably evaluating the operation the result of the clone is not a node but a tensor variable.

(Pdb) theano.printing.debugprint(dependent_rv.logp_elemwiset.clone())
<TensorType(float32, vector)> [id A]

pymc3_hmm/distributions.py

pymc3_hmm/step_methods.py

xjing76 · 2021-03-26T21:04:21Z

Here I get the traversing function to be conditioned on SwitchingProcess. And get the location of the alloc without traversing.

However, I think I am still failing the FFB unit test. The states that we are getting is not the state that we suppose to get. If @brandonwillard you can review this portion that would be great, as I am not confident that it is the right way to set the S_t. What I am trying to do is to set the St to be all 0 and all 1 for the two states just like in the original setup. However, the seems like for every run of ffb.step() the result for states are different, which makes me questions if this is right.

-> assert np.array_equal(res["S_t"], poiszero_sim["S_t"])
(Pdb) poiszero_sim["S_t"]
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 1, 1, 0, 0, 0, 0, 0])
(Pdb) ffbs.step(test_point)['S_t']
array([0, 0, 1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 0, 0, 1,
       0, 0, 0, 1, 1, 1, 0, 1])
(Pdb) ffbs.step(test_point)['S_t']
array([0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0,
       0, 0, 1, 1, 0, 0, 0, 1])
(Pdb) ffbs.step(test_point)['S_t']
array([1, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0,
       1, 0, 0, 1, 0, 0, 1, 1])

pymc3_hmm/step_methods.py

brandonwillard · 2021-03-30T20:08:46Z

pymc3_hmm/step_methods.py

+        node = queue.pop(0)
+        if node not in visited:
+            visited.append(node)
+            if node.__str__() =="Alloc.0":


Use isinstance on node.owner.op instead.

xjing76 force-pushed the ffb branch from 697613b to 112f820 Compare March 16, 2021 22:35

brandonwillard suggested changes Mar 18, 2021

View reviewed changes

brandonwillard changed the title ~~Trying to improve FFB~~ Reuse allocated arrays in FFBS sampler Mar 18, 2021

brandonwillard marked this pull request as draft March 18, 2021 01:00

xjing76 force-pushed the ffb branch 2 times, most recently from a2b711e to f53dcc5 Compare March 18, 2021 03:08

xjing76 force-pushed the ffb branch 3 times, most recently from 9fd34e4 to f9abf35 Compare March 22, 2021 18:09

xjing76 force-pushed the ffb branch 4 times, most recently from efa3bfb to b447b71 Compare March 25, 2021 16:32

brandonwillard suggested changes Mar 26, 2021

View reviewed changes

xjing76 force-pushed the ffb branch 2 times, most recently from 82f6892 to d76747e Compare March 26, 2021 20:47

xjing76 requested a review from brandonwillard March 26, 2021 21:33

brandonwillard suggested changes Mar 30, 2021

View reviewed changes

Fix dtype resolution in SwitchingProcess

25efe93

brandonwillard force-pushed the ffb branch 2 times, most recently from dce65ce to d26de7e Compare March 31, 2021 03:37

Use BlockedStep base class for FFBSStep and update output in-place

ec8dfdf

brandonwillard force-pushed the ffb branch from d26de7e to ec8dfdf Compare March 31, 2021 03:49

brandonwillard marked this pull request as ready for review March 31, 2021 03:49

brandonwillard merged commit 4420feb into AmpersandTV:main Mar 31, 2021

brandonwillard mentioned this pull request Mar 31, 2021

FFBS performance improvements #76

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reuse allocated arrays in FFBS sampler #81

Reuse allocated arrays in FFBS sampler #81

xjing76 commented Mar 16, 2021 •

edited

Loading

brandonwillard Mar 18, 2021

xjing76 commented Mar 19, 2021 •

edited

Loading

xjing76 commented Mar 22, 2021

xjing76 commented Mar 23, 2021 •

edited

Loading

xjing76 commented Mar 26, 2021

brandonwillard Mar 30, 2021

Reuse allocated arrays in FFBS sampler #81

Reuse allocated arrays in FFBS sampler #81

Conversation

xjing76 commented Mar 16, 2021 • edited Loading

brandonwillard Mar 18, 2021

Choose a reason for hiding this comment

xjing76 commented Mar 19, 2021 • edited Loading

xjing76 commented Mar 22, 2021

xjing76 commented Mar 23, 2021 • edited Loading

xjing76 commented Mar 26, 2021

brandonwillard Mar 30, 2021

Choose a reason for hiding this comment

xjing76 commented Mar 16, 2021 •

edited

Loading

xjing76 commented Mar 19, 2021 •

edited

Loading

xjing76 commented Mar 23, 2021 •

edited

Loading