Docker GPU image for PyMC3

I’m interested in creating a Docker image that can use the GPU for PyMC3. Currently, I am working on combining this image with the one listed in the PyMC3 github repository.

Does anyone have any suggestions or best practices? Advice regarding combinations of CUDA, CuDNN, and Theano/Theano-PyMC3/Aesara are especially appreciated!

1 Like

Very cool, definitely post once you have something you like.

I would perhaps focus on the JAX backend where we saw great speed-ups on the GPU.

1 Like

Hi ckrapu, did you get anywhere with this? I’m hitting various walls trying to get this working.

For ref, this is where i’ve got to. GitHub - alephinsights/docker-pymc-gpu

IIRC, I think I did get it to work. I can’t share my entire Dockerfile but the hardest part was getting the right versions lined up. Here were the package versions and base image that I used.

FROM nvidia/cuda:8.0-cudnn6-devel

thanks for the steer. still little luck.

were you installing via conda or pip? Did you have any luck with more recent cuda/cudnn?