Question about notebook The prevalence of malaria in the Gambia

Questions v5
1

I am trying run this example notebook : The prevalence of malaria in the Gambia — PyMC example gallery

My environment is Jupyter notebook running in Windows 11 Pro with GPU

pymc version 5.27.1

I am running into an issue at the section: Out-of-sample posterior predictions

The following step is taking a very long time: the progress bar hasn’t moved from 0% for about 1 hr

with hsgp_model:
pm.set_data(new_data={“X”: new_lonlat, “elevation”: new_elev_std})
pp = pm.sample_posterior_predictive(hsgp_trace, var_names=[“p”], random_seed=rng)

Is this normal? Any way to speed this up? I have not found any pointer.

Thanks!

2

Can you try to run on google colab and see if it also hangs there?

3

Thanks for your response.

I restarted jupyter kernel and then ran the notebook again. That process was not “hanging” - the sampling progress bar was not refreshing. The step completed in 29 minutes.

As you suggested I ran the notebook in colab today and the step completed in about 29 minutes - also, the sampling progress bar refreshed to show the progress.

Any way to speed up this step?

Thanks!!

4

You can try a different compute backend by adding compile_kwargs={'mode':'JAX'} or {'mode':'NUMBA'} to the pm.sample_posterior_predictive call. Since the model is just a big matmul, if you have a GPU you might be able to get some lift in jax mode. Numba should be about the same since it’s all BLAS calls anyway.

You can also use a coarser grid – the notebook uses new_lonlat = elevations[["y", "x"]].to_numpy() which is probably huge – you could reduce the number of points you’re sampling.

1 Like
5

OK..I will try running it in JAX mode, coarse grid.

How to enable GPU mode or is GPU detected and enabled automatically?

Thanks!!

6

jax mode will do whatever you have jax configured to do

7

Thanks!!

8

Did using JAX give you anything? Or did you just end up using a finer grid?

9

JAX didn’t help much - perhaps because I didn’t install JAX with cuda –> This will be another effort navigating requirements for Windows 11 (pro) - I don’t think this possible in Windows. Coarser grid did help. NUMBA didn’t help either (as I had expected).

So, best option for windows 11 is use coarser grid.

Thanks

1 Like