Tensorflow with GPU in docker on Google Compute Engine
First install these on a new VM:
docker
Installation guide:
https://docs.docker.com/engine/install/ubuntu/docker-compose(optional)
Installation guide:
https://docs.docker.com/compose/install/nvidia-docker2
Installation guide:
https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html
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
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.pyThen you can simply run docker-compose up and start your training/prediction job. The runtime:nvidia ensures that docker can see our GPU.