Google Colab is a Jupyter Notebook-based cloud service which gives free GPU use. I am wondering if anyone of you managed to get a working PyMC3 environment with GPU support. I struggle to get all the dependencies right (CUDA, pyGPU, etc.). I have a tricky model, and I would like to see if the GPU speeds things up a bit.
+1
Do you have a colab to show the things you had already tried?
Here’s the ipynb notebook Theano_with_GPU.txt (401.0 KB) (please, change the extension to .ipynb, otherwise I could not upload the file).
It is sort of working now. The Theano test shows that the GPU is in use, but the sampling is much slower compared to using the CPU only (we are talking about 1297 draws/sec with CPU compared to 270 draws/sec with GPU). PyGPU still complains that the CUDNN is not setup correctly, so maybe that’s why it is so slow. Maybe you can start from this example. I also noticed that the PyMC3 works only with 1 chain when the GPU is in use.
I started with the assumption that Colab fullfilled all the CUDA requirements. At first, I tried to patch all the errors that PyGPU was raising, trying to find the missing files in the system directories. But then more errors were raised. Therefore, I ended up reinstalling CUDA from scratch. The thing that is missing is a cudnn.h file that I cannot find.
1 Like
Regardless of the colab setup: as far as I know/have read most PyMC3 models generally do not see an increase in (sampling) speed when using a GPU instead of CPU.
My experience is that, when using the GPU, I need to set cores=1 when sampling.
See also:
I also tried colab unsuccessfully. It seems that it doesn’t support theano! Potentially because we can’t setup the theano flags.
1 Like
Is there any update on this topic since 2018? Does anyone know if it’s possible to use PyMC3 with GPU on Google Colab?
Have you try PyMC v4 with the Jax backend? should work out of the box.
No I haven’t, do you have any example code? I struggle to find these things clearly explained somewhere. A vignette or an actual applied example could be very helpful.
We are in the process of updating the doc, meanwhile this blog post (MCMC for big datasets: faster sampling with JAX and the GPU - PyMC Labs) should help (code here)
cc @Martin_Ingram in case there is additional question.
3 Likes
Thanks. Also I am not sure how to install PyMC v4 on Google Colab. Is it straightforward?