Running TensorFlow with Docker and GPU

Georgy Georgiev
2 min readApr 18, 2016

--

For those who had trouble getting Google’s TensorFlow to work inside a Docker container while accessing NVIDIA GPU-s with CUDA 8.0 and CUDNN 7 in the past months, these next lines are for you.

1. Install CUDA on host

First you need to install Cuda 8.0 on your host. Detailed instructions of how todo it on a Ubuntu 14.04 are available here. Cuda provides the necessary drivers and libs for running computations on NVIDIA GPU-s.

Make sure to set the LD_LIBRARY_PATH and CUDA_HOME environment variables. Consider adding the commands below to your ~/.bash_profile. These assume your CUDA installation is in /usr/local/cuda:

export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64"
export CUDA_HOME=/usr/local/cuda

2. Install CUDNN on host

At the time of the writing (Sep 2017) it is best to use Cudnn v7. Register at nvidia.com and download the install archive.

tar xvzf cudnn-8.0-linux-x64-v7.tgzsudo cp cuda/include/cudnn.h /usr/local/cuda/includesudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64sudo chmod a+r /usr/local/cuda/lib64/libcudnn*

3. Install Docker and nvidia-docker

# Install docker
curl -sSL https://get.docker.com/ | sh
The docker container needs access to the GPU devices. For this purpose use `nvidia-docker` which is a wrapper around the standard `docker` command.# Install nvidia-docker and nvidia-docker-pluginwget -P /tmp https://github.com/NVIDIA/nvidia-docker/releases/download/v1.0.0-rc.3/nvidia-docker_1.0.0.rc.3-1_amd64.debsudo dpkg -i /tmp/nvidia-docker*.deb && rm /tmp/nvidia-docker*.deb# Test nvidia-smi.
nvidia-docker run --rm nvidia/cuda nvidia-smi

You might need to use `nvidia-docker` with sudo!

4. Run a Tensorflow GPU-enable Docker container

The container itself is started as pointed out in the official documentation as follows:

# Run container
nvidia-docker run -d --name <some name> -p 8888:8888 -p 6006:6006 gcr.io/tensorflow/tensorflow:latest-gpu
# Log in
nvidia-docker exec -it <some name> bash

e.g.:

nvidia-docker run -d --name tf1 -p 8888:8888 -p 6006:6006 gcr.io/tensorflow/tensorflow:latest-gpunvidia-docker exec -it tf1 bash

Note: Port 8888 is for ipython notebooks and port 6006 is for TensorBoard.

You can test if everything is alright by running this Python script.

If you have comments/questions, you can also drop me a line at georgy@vnm.io or by using my Twitter handle @georgy_georgiev.

--

--