PyMC3 MCMC performance with and without Theano's NumPy BLAS warning (updated with PyMC v4 comparison!)

Introduction

#

PyMC3 is a Python-based probabilistic programming language used to fit Bayesian models with a variety of cutting-edge algorithms including NUTS MCMC¹ and ADVI². It is not uncommon for PyMC3 users to receive the following warning:

WARNING (theano.tensor.blas): Using NumPy C-API based implementation
for BLAS functions.

where Theano³ is the autodiff engine that PyMC3⁴ uses under-the-hood. The usual solution is to re-install the library and its dependencies following the operating system-specific instructions in the wiki.

I am working on a project using PyMC3 to fit large hierarchical models and was receiving the warning on the Linux server used to fit the full models, but not on my local MacBook Pro when experimenting with model designs (even after re-creating the conda environment). Therefore, I asked for more help on the PyMC Discourse with two questions:

1. How can I get this warning to go away?
2. Is the problem that Theano is warning about affecting its performance?

Thankfully, I got the first question squared away (it was a problem with my $PATH) and that enabled me to test the second. The results of that test are presented here.

Update

#

Based on a suggestion from PyMC core developer, Thomas Wiecki, on Twitter, I have since updated this analysis by fitting the same model on the same data on the same compute cluster with PyMC v4 (specifically v4.0.0b2). I left the text below primarily untouched, but the main take away was that Wiecki was correct (unsurprisingly): version 4 is much faster than version 3.

TL;DR

#

1. Does Theano’s warning “Using NumPy C-API based implementation for BLAS functions,” affect its performance?
2. On average, sampling was faster without the warning.
3. But, the rate of sampling at its fastest point was the same in each condition.
4. Thus, fixing the underlying issue for the warning may increase sampling, but there is still a lot of variability due to the stochastic nature of MCMC.
5. MCMC sampling PyMC v4 (currently in beta) is faster than PyMC3.

The experiment

#

For my test, I fit a large hierarchical model I am working on with and without the amended $PATH variable, the latter resulting in the Theano warning. The model is for a research project that is still hush-hush so I can’t publish the code for it just yet, but it is a negative binomial hierarchical model with thousands of parameters. The model was fit with just under 1 million data points. I fit four chains for each round of MCMC, each with a different random seed and 1,000 draws for tuning and 1,000 draws for sampling. Each chain was run on a separate CPU with up to 64 GB of RAM. Using a custom callback, I was able to log the time-point for every 5 draws from the posterior.

Results

#

Resource consumption

#

To begin, each model used about 53-56 GB of the available 64 GB of RAM. No real substantial difference there.

Sampling rates

#

The first plot below shows the time-course for each chain, colored by the experimental condition. The plot below that displays the duration of each 5-draw interval indicating the rate of the sampling process over time for each condition followed by a plot that shows the same data on the same axes.

Each chain generally went through about 3 stages of rapid early tuning, slow tuning, and then rapid sampling post-tuning (the draws that will represent the approximate posterior distributions).

One exception to this is one of the chains from the “warning” condition that, comparatively, had an incredibly rapid early stage, prolonged slow stage, and a slower final stage than the other chains. (This chain took longer than 5 days to fit and my job timed-out before it could.) I similar result happens reliably when I use ADVI to initialize the chains (this may be a future post), so I think this is just a process of the randomness inherent to MCMC and not necessarily attributable to the Theano warning. Removing this chain shows how similar the results were between those remaining.

The sampling durations in each condition are also plotted as histograms (below).

Summary table

#

Finally, the following is a table summarizing the sampling rates for each chain. Note that chain #2 of the “warning” condition never finished (>5 days).

Note that the average draw rate (in minutes per 5 draws) is the same for all chains during the “sampling” stage (except for the outlier chain).

Conclusion

#

My understanding of the inner-workings of PyMC and Theano is limited, so it is impossible for me to provide a confident final conclusion. From this experiment, it seems like PyMC3 performed equivalently with and without the warning message. That said, it is probably best to address the underlying issue to ensure optimal PyMC3 performance and behavior.

Update with addition of PyMC v4: PyMC v4 samples faster during both the tuning and sampling stages of MCMC. In addition, version 4 tends to have fewer long-duration steps in the tuning process, dramatically cutting down on overall runtime.

Featured image source: “Theano – A Woman Who Ruled the Pythagoras School”, Ancient Origins.