Setting up TensorFlow with GPU acceleration the quick way

Photo by Will Truettner on Unsplash

Introduction

TensorFlow, with its 2.x release this year, including changes such as the “folding in” of Keras, seems to have “won” as the dominant Deep Learning library. With the aforementioned additions, the library has become easier to use than ever. However, setting it up, especially in order to take advantage of GPU acceleration (practically required for any non-trivial project) remains something of a chore.

This post aims to highlight a specific way to set up your TensorFlow training runs that obviates the need for the somewhat cumbersome installation process. We’ll do that by employing containerization, which will isolate the required packages to the container image, and not mess with our OS.

Setting up Docker

Let’s now focus on the containerization process itself — both the basics, and some customization.

nvidia-docker

Assuming we have an NVIDIA GPU (some sort of support for AMD/Radeon GPUs also exists, but wasn’t tested here), and a Linux machine available (macOS and Windows are unsupported sadly), we can use NVIDIA’s specialized Docker runtime. To do so, simply follow the installation instructions for your distribution.

Once we’re through with that, we should be able to launch the test script container with either:

or

(the latter is the old syntax variant; you may also encounter examples with nvidia-docker as the command - those are obsolete)

If a list of GPUs appears corresponding to the output of nvidia-smi outside of the container, we’re all set for the next step.

Creating and launching the TensorFlow container

Of course, we’re not really interested in the output of GPU driver diagnostics, but in performing actual, GPU-assisted, TensorFlow runs. For that, we turn to the official TensorFlow Docker images.

Configuration of the image variants is arranged through tag suffixes. We’re mostly interested in the -gpu ones, but if you want to have an integrated Jupyter notebook server, you can also try the -jupyter meta tags. Note that tags compose, so to get Jupyter with GPU support, you should use the tag <VERSION>-gpu-jupyter, e.g. latest-gpu-jupyter.

Let’s perform our first test run:

This should open an interactive (--it) Python 3 shell in a disposable (--rm) container.

NOTE: If you’ve been using the image by any chance before April, you need to execute docker pull tensorflow/tensorflow:latest-gpu to get the Python 3 shell, due to the Python 2 EOL Changes. This is also why there’s no py3 suffix for image labels now.

Now, let’s check the version of TensorFlow:

(current at the time of writing)

And let’s test the GPUs:

In most cases concerning “desktop” machine learning development, you will want to enable memory growth as well:

As a final step, you might like to map a local directory (the one you have your code in presumably) with a volume on the container. To do that, just use the standard -v option:

Note that we’re mapping to a subdirectory of /tf, as this is the default working dir for the image - especially relevant when working with Jupyter.

Rolling your own custom container

We’ve got the basics settled now. However, when developing machine learning processes and pipelines, we usually require more than just TensorFlow. That means more packages needing to be installed.

We could either run a pip install <package> through docker exec on the container, or use the Jupyter variant and execute !pip install <package> in a notebook. However, assuming we want to keep the "disposability" advantage of the containers, we’d need to run those commands every time we start a given run.

The alternative here is to roll our own container image, based on the TensorFlow one.

First, let’s create a new directory from which we’ll build our image from. In this directory, we’ll add a requirements.txt for our additional dependencies, for example:

Now let’s add a Dockerfile that defines our build:

Finally, let’s run a build command in the created directory:

This will result in an <your image name>/tensorflow-gpu-jupyter installed in your local Docker cache, available for use.

NOTE: You can also push this image to a registry of your choosing, thereby making it available to your colleagues.

Testing out the image

Now that we’ve set up our container, it’s time to try it out. We’re going to use the MNIST tutorial for the test, with some slight modifications.

Let’s start with running a fresh instance of the container. We’ll add a port mapping to the host machine for simplicity.

Note that:

• you can also substitute your custom image for this run;
• we are not mapping any volume here, since everything that we need is either already on the container, or can be downloaded.

We should now see something like this:

Copy the provided URL into your browser. This should result in Jupyter’s notebook interface showing up:

As you have undoubtedly noticed, a tutorial directory is included in the image by standard.

Let’s open up a new notebook, and fire up the same test code as before:

The same list of GPUs as in the console should appear.

Now, let’s go back to the tutorial in the subsequent cells. But before we do that, we need to install the tensorflow_datasets package. If you’ve followed the optional section and are using your custom container, you already have it. Otherwise, execute, preferably in a fresh cell:

Subsequently:

If you got the standard training-by-epoch output, you’re all good:

Conclusions

This blog entry demonstrates an alternative, arguably more reliable and/or accessible, way to use GPU acceleration with TensorFlow. It also details tailoring a custom container image to your needs, based on the “official” TensorFlow ones.

Keep an eye out for a follow-up post, showing how the same thing can be done while also taking advantage of a proper IDE.