Photo by Caspar Camille Rubin on Unsplash

Hi folks,

In this article, I want to share with you a very short and simple way how to use Nvidia GPU in docker to run TensorFlow for your machine learning (and not only ML) projects.

Prerequisites:

Nvidia videocard
Ubuntu 18.04
Installed Nvidia driver

Let's go

Install Docker on your OS:

sudo apt-get install docker-ce
docker run hello-world #Test docker that it installed correct:

Install official Nvidia docker:

Add Nvidia repository to get Nvidia docker container (these code can be executed on various Ubuntu versions):

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 container with gpu support and restart docker daemon:

sudo apt-get install nvidia-docker2
sudo pkill -SIGHUP dockerd

After that, you have Nvidia runtime with support of your gpu in a container! Remember that you can use it with any Docker container.

Now we can check our nvidia runtime by calling nvidia-smi (nvidia-smi tool used for monitoring your GPU):

docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi

Profit! :)

Let's bench our configuration

To do it we will use script:

import sys
import numpy as np
import tensorflow as tf
from datetime import datetimedevice_name = sys.argv[1]  # Choose device from cmd line. Options: gpu or cpu
shape = (int(sys.argv[2]), int(sys.argv[2]))
if device_name == "gpu":
    device_name = "/gpu:0"
else:
    device_name = "/cpu:0"with tf.device(device_name):
    random_matrix = tf.random_uniform(shape=shape, minval=0, maxval=1)
    dot_operation = tf.matmul(random_matrix, tf.transpose(random_matrix))
    sum_operation = tf.reduce_sum(dot_operation)startTime = datetime.now()
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as session:
        result = session.run(sum_operation)
        print(result)# It can be hard to see the results on the terminal with lots of output -- add some newlines to improve readability.
print("\n" * 5)
print("Shape:", shape, "Device:", device_name)
print("Time taken:", str(datetime.now() - startTime))

Let's start with using original tensorflow image:

docker run --runtime=nvidia --rm -ti -v "${PWD}:/app" tensorflow/tensorflow:latest-gpu-jupyter python /app/benchmark.py cpu 10000
docker run --runtime=nvidia --rm -ti -v "${PWD}:/app" tensorflow/tensorflow:latest-gpu-jupyter python /app/benchmark.py gpu 10000

Results:

# with CPU (docker)
(‘Shape:’, (10000, 10000), ‘Device:’, ‘/cpu:0’)
(‘Time taken:’, ‘0:00:02.699631’)# with GPU (docker)
(‘Shape:’, (10000, 10000), ‘Device:’, ‘/gpu:0’)
(‘Time taken:’, ‘0:00:00.804804’)

To compare native start on Ubuntu 18.04 (without docker):

Shape: (10000, 10000) Device: /cpu:0
Time taken: 0:00:04.583930Shape: (10000, 10000) Device: /gpu:0
Time taken: 0:00:00.783128

Conclusion

As a result, we can use our Nvidia GPU in a docker container. Our benchmark result shows us that we gain huge speed up when we use gpu solution in docker. Moreover, I got better results in the docker container than on the native OS on my configuration.

Used materials:

https://marmelab.com/blog/2018/03/21/using-nvidia-gpu-within-docker-container.html

https://nvidia.github.io/nvidia-docker/

https://medium.com/@sh.tsang/docker-tutorial-5-nvidia-docker-2-0-installation-in-ubuntu-18-04-cb80f17cac65

#Nvidia:

How to use Nvidia GPU in docker to run TensorFlow