How to install Docker for GPU, install a Tensorflow GPU Docker container and Run a a Jupyter Notebook for Machine Learning

Part 1: Checking the prerequisites

In this example I will explain how to install Docker for GPU and a Docker container with TensorFlow on Ubuntu 20.04. This assumes you have already installed the GPU driver with the correct version of CUDA for the given TensorFlow version. Currently I have installed TensorFlow 2.5 with CUDA version 11.2. TensorFlow is not compatible with newer CUDA versions such as 11.4, nor are newer graphics drivers compatible with older CUDA versions, so make sure you check here before you choose which driver to install. You will probably need to install an older CUDA version and graphics driver than the latest available to support the latest TensorFlow version.

Part 2: Installing Docker for GPU

You must install a special version of Docker to work with the GPU. I found the details on setup and installation here. Open a terminal and enter the following command to run Docker’s convenience script to point Ubuntu to Docker’s installation files.

curl | sh \
  && sudo systemctl --now enable docker

Now enter this command to point to the specific installation files for GPU-compatible Docker.

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
   && curl -s -L | sudo apt-key add - \
   && curl -s -L$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

Now update Ubuntu’s repositories so all the new information you pointed to is readily available.

sudo apt-get update

Now you are able to install the nvidia GPU compatible Docker version

 sudo apt-get install -y nvidia-docker2

Next, restart Docker to ensure the installation changes take effect.

sudo systemctl restart docker

Now you can enter the command below to make sure Docker can run the nvidia-smi to indicate if everything worked correctly.

sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

In this case, for TensorFlow 2.5, you should get the following output with CUDA version 11.2:

Success! Now we can install a Docker container pre-loaded with TensorFlow.

Part 3: Installing Docker Container with Tensorflow and Running Examples

You only need to enter one line into the terminal to install and run the TensorFlow docker container, then run the Jupyter notebook from within that container:

docker run --gpus all -it -p 8888:8888 tensorflow/tensorflow:2.5.0-gpu-jupyter --no-browser

The TensorFlow install page has some details on how to choose the TensorFlow version in your command. In this case, I wanted to make sure to use TensorFlow version 2.5.0 because it is relatively new, compatible and stable with my current CUDA installation. The “-p 8888:8888″ and ” –no-browser” commands allow Jupyter to be accessed from our computer which is  outside the Docker container environment. The first time you run the command, it may take longer because Docker will automatically detect if you have the relevant container and install it if needed.

Once the command is finished running, you should see a link somewhere in your terminal. You can press “CTRL” and click the link to open the TensorFlow tutorials directory into your browser.

Now you can click the “tensorflow-tutorials” link and choose any of the example “.ipynb” files to run. In this case, I chose the “classification.ipynb” example. Simply click the example, go to the page and click the double forward to run through the entire example at rapid pace using your GPU.

Congratulations! Now you can read through the examples, see how they work, and start learning Tensorflow at a fast pace!

Leave a Reply