Tensorflow with GPU in docker on Google Compute Engine

First install these on a new VM:

nvidia-docker2 is used to make the GPU available to our container Make sure you get a GPU for your VM.

Getting CUDA and CuDNN and TensorFlow to live together in harmony

I remember I used to try and figure out which version of python, tensorflow, GPU driver, CUDA driver, cuDNN work together, repeating apt remove --purge '^nvidia-.\*' for a zillion times to no avail. Here comes docker for our help.

Option 1

We can use docker images provided by Nvidia, they come in many flavors (Ubuntu, CentOS, etc.) and it installs the proper CUDA and cuDNN libraries.

You can find a list of available docker images on their docker hub site: https://hub.docker.com/r/nvidia/cuda

Example: https://gitlab.com/nvidia/container-images/cuda/blob/master/dist/11.2.1/ubuntu20.04-x86_64/devel/cudnn8/Dockerfile

The above one comes with CUDA11.2, CuDNN8 installed in Ubuntu20.04. Pretty neat!

Option 2 - The easier one!

We can use tensorflow-gpu as base image out of the box. It comes installed with everything necessary. https://hub.docker.com/r/tensorflow/tensorflow

For this post, I will use tensorflow/tensorflow:latest-gpu docker image.

The only thing we need to make sure is to use nvida-docker when running the image.

Dockerfile example

FROM tensorflow/tensorflow:latest-gpu
RUN mkdir /app
WORKDIR /app
ADD . /app/
RUN pip install --upgrade pip
RUN pip install -r requirements.txt

Build your image and run:

docker build -t {YOUR-TAG} .
docker run --gpus all -it {YOUR-TAG} python main.py

That’s it! Notice how we passed the --gpus all option to docker run, this is the magic parameter that tells docker to use our nvidia-docker2 runtime.

docker-compose example

For all our docker-compose aficionados out there, here is a docker-compose.yaml sample:

version: "3"
services:
you-service-name:
runtime: nvidia
build: .
volumes: - ./model:/app/model
command: python main.py

Next you can simply run docker-compose up and start your training/prediction job. The runtime:nvidia part ensures that docker can see our GPU.