Setting up your GPU for Tensorflow 2.4 (2021)

Setting up your GPU with Tensorflow does not have to be a nightmare. Image: nvidia.com

TensorFlow 2.4 is out now. Which means that you probably won’t update because you are afraid of going through the whole setting your GPU nightmare once again. You are probably using CUDA 10.1 and it just works, so why bother right? Well I have great news for you: setting up TensorFlow 2.4 on Windows 10 with full CUDA 11.0 support is not that hard. Here I will provide a full guide on how to setup your GPU in no time (unless you are me and takes you 2 days because you could not find a good tutorial).

First, we need to understand that both CUDA and Tensorflow are dynamic, they update their releases and things might get broken once in a while. That is why is so important, and I can stress this enough, to match the correct version of TF and CUDA (if you stop reading at least take that message home).

System requirements from the TensorFlow.org (Jan 2021)

Based on those specification and previous experiences, I successfully used the following steps:

1. Update the NVIDIA drivers to the latest version

First things first: use your GPU to its full potential and update the drivers. For this tutorial I am using a Razer 15 Advanced with RTX 2080 Super Max Q. For updating to the most recent version, just go to the NVIDIA website and look for the latest drivers for your GPU. Restart your computer afterwards (please just restart as much as possible, because you know, Windows).

2. Install the CUDA toolkit

After updating the NVIDIA drivers we can install the CUDA toolkit which includes multiple libraries that tensorflow requires on loading. The most common error when trying to run TensorFlow with GPU is a complaint that CUDA .dll files could not be found. The Nvidia experience app gives you the latest version for your GPU card or just can just download manually from the drivers website.

Important: For TensorFlow 2.4 you must download CUDA 11.0 (not 11.1 or 11.2, just don’t).

cuDNN 8.0.4 works with TF 2.4 and CUDA 11.0 (aka as the killer combo).

Since cuDNN does not really require installation, you have two options:

  1. Add the cuDNN to the Windows PATH (recommended)
  2. Copy the cuDNN files to the corresponding folders in the CUDA 11.0 installation folder.

To be on the safe side let’s add all environment variables to the Windows path.

Here we assume that you copied the contents of cuDNN to C:\tools\cuda

SET PATH=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.0\\bin;%PATH%SET PATH=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.0\\extras\\CUPTI\\lib64;%PATH%SET PATH=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.0\\include;%PATH%SET PATH=C:\\tools\\cuda\\bin;%PATH%

3. Install Tensorflow 2.4 with Anaconda

Finally, we can get to business. Anaconda is probably the best way for using python in Windows. So let’s install Anaconda (or miniconda if you are resource limited). After installing, open the Anaconda terminal and create a new environment:

# check your python version! 
conda create -c conda-forge python=3.8.5 -n tf
conda activate tf

Create a new environment and install TF

pip install tensorflow 
# will automatically install TF 2.4

Open python on the Anaconda terminal and run:

import tensorflow as tf 
print(tf.reduce_sum(tf.random.normal([1000,1000])))

The output should tell you that all libraries were loaded successfully!

And that’s that! Now you can go and train some beautiful deep learning models using your NVIDIA GPU.

Computational biologist and data scientist. Data narrative on molecular biology, machine learning and finance.