I found that the easiest and hassle-free way of running the GPU version of Tensorflow is with the help of docker as you don’t have to worry about the CUDA versions and Tensorflow versions compatibility. This is a tutorial to use Tensorflow in a docker container in Ubuntu. If you were to use Tensorflow in a docker container in Mac or Windows, it will only be simpler than this tutorial because you just need to install Docker for Desktop application for Mac or Windows as opposed to a Docker Engine.
Install Docker CE in Ubuntu
Uninstall old versions
sudo apt-get remove docker docker-engine docker.io containerd runc
Set up the repository
sudo apt-get install \
apt-transport-https \
ca-certificates \
curl \
gnupg-agent \
software-properties-common
Add Docker’s official GPG key
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo apt-key add -
Verify the fingerprint
sudo apt-key fingerprint 0EBFCD88
Add the repository
sudo add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable"
sudo apt-get update
sudo apt-get install docker-ce docker-ce-cli containerd.io
Verify the Docker CE installation
sudo docker run hello-world
Install nvidia-docker2 for GPU support (Only necessary if you’re planning on running Tensorflow in GPU)
If you have nvidia-docker 1.0 installed: we need to remove it and all existing GPU containers docker volume
ls -q -f driver=nvidia-docker | xargs -r -I{} -n1 docker ps -q -a -f volume={} | xargs -r docker rm -f sudo apt-get purge -y nvidia-docker
Add the package repositories
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey|sudo apt-key add -
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
Install nvidia-docker2 and reload the Docker daemon configuration
sudo apt-get install -y nvidia-docker2
sudo pkill -SIGHUP dockerd
Test nvidia-smi with the latest official CUDA image
docker run --runtime=nvidia --rm nvidia/cuda:9.0-base nvidia-smi
Get Tensorflow docker image from docker hub
# run one of the following commands based on your need. Tag variants are
-gpu, -py3 and -jupyter
sudo docker pull tensorflow/tensorflow:latest-gpu-py3 #pulls the latest tensorflow docker image with GPU support and python3
sudo docker pull tensorflow/tensorflow:latest #pulls the latest tensorflow docker image with CPU only support and python2
Instantiate your container from your tensorflow docker image
# for tensorflow CPU image with python 2
sudo docker run -it --name=mytfcontainer tensorflow/tensorflow:latest bash
# for tensorflow GPU image with python 3
sudo docker run -it --runtime=nvidia --name=diwanshu1 tensorflow/tensorflow:latest-gpu-py3 bash
There you have it. Enjoy using tensorflow in docker
References: