PyMC3 MCMC performance with and without Theano’s NumPy BLAS warning

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.

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.

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.

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.

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

1. Hoffman, Matthew D., and Andrew Gelman. 2011. “The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo.” arXiv [stat.CO]. arXiv. http://arxiv.org/abs/1111.4246. ↩︎

2. Kucukelbir, Alp, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M. Blei. 2017. “Automatic Differentiation Variational Inference.” Journal of Machine Learning Research: JMLR 18 (14): 1–45. ↩︎

3. Theano was the wife (and follower) of the famous mathematician Pythagoras and Aesara was their daughter and followed in their occupational and religious footsteps. ↩︎

4. The next version of PyMC3 will be called “PyMC” and use a fork of Theano called Aesara. ↩︎