Installing Tensorflow with GPU support on Linux

  • Machine Learning

Lately I spent some time using Deep Learning and configuring Linux servers with GPUs so the models train faster. In this short blog post I am going to list all what you have to follow in order to properly install the nvidia drivers, cuda, and other tools you'll need before running Tensorflow (or Keras) with GPU support. Let's get to it !

Note that this tutorial can work either on Ubuntu or Debian. This tutorial will install tensorflow-gpu 1.12 which is the latest version available while I am writing these words.

Tensorflow GPU needs the following softwares:

  • NVIDIA® GPU drivers (CUDA 9.0 requires 384.x or higher)
  • CUDA® Toolkit (TensorFlow supports CUDA 9.0)
  • cuDNN SDK (>= 7.2)

Installing on Ubuntu

Here is a complete shell script showing the different steps to install tensorflow-gpu:

Docker Image

If you are familiar with Docker, I recommend you have a look at the Tensorflow Docker Image. It’s already configured with the latest drivers and can run on CPU or GPU.

If you’re not familiar with Docker, you should definitely learn using it. Here are some links to get you started:

Install Keras on top of Tensorflow

This step is not required, and some people probably prefer using Tensorflow directly with no abstraction layer on top of it. I personally started doing Deep Learning with Keras on top of Tensorflow, because it provided a simpler API, and I find it really easy and fast to build models. Keras was build by François Chollet, and since he is now working for Google, Keras is very well integrated with Tensorflow.

Installing keras is as easy as pip install keras. It will automatically detect your GPUs if you have tensorflow-gpu installed, like we did.

To get started, have a look at the official Keras website and their getting started guide. I have also wrote a blog post with some Keras snippets that I find useful.

I hope this post was helpful, and have fun with Deep Learning !