CUDA enabled Jupyter Docker Images

I am happy to announce that Jupyter Docker Stacks project now provides GPU accelerated Docker images. This allows PyTorch or TensorFlow operations to use compatible NVIDIA GPUs for accelerated computation.

How to find these images

All our images are available on Quay.io registry. We provide CUDA accelerated versions of images are available for tensorflow-notebook and pytorch-notebook.

To use such an image, you have to specify a special prefix tag to the image: versioned CUDA prefix like cuda11- or cuda12- for pytorch-notebook or just cuda- for tensorflow-notebook.

We build pytorch-notebook only for 2 last major versions of CUDA, tensorflow-notebook image supports only the latest CUDA version listed in the officially tested build configurations list. CUDA-enabled images are available on x86_64 platform.

For example, you can use an image quay.io/jupyter/pytorch-notebook:cuda12-python-3.11.8 or quay.io/jupyter/tensorflow-notebook:cuda-latest.

Running the images

To enable CUDA support on Linux host machine, you need to:

• have compatible Nvidia GPU
• have NVIDIA Linux driver installed
• add --gpus all (or --gpus '"device=all"') flag to if you’re using Docker
• add --device 'nvidia.com/gpu=all' flag if you’re using Podman

For example, you can run the image like this: docker run -it --rm -p 8888:8888 --gpus all quay.io/jupyter/tensorflow-notebook:cuda-latest.

You can also enable GPU support on Windows using Docker or Podman.

Reaching out

You will find our documentation here.
If you encounter any problems, please, take a look at GitHub issues.
If you would like to improve GPU support in our images, please, reach out here.

Acknowledgments

I would like to thank:

• Johanna Reiml, for implementing general variant concept for our images and making it work for pytorch-notebook (PR #2091)
• Christof Kaufmann, for making it work for tensorflow-notebook (PR #2100)
• Olivier Benz, for helping me out to write this blog post