Project maintained by markjay4k Hosted on GitHub Pages — Theme by mattgraham

Install Tensorflow 2.4

Github pages version

We’re going for Tensorflow 2.4.1 with GPU support on an Ubuntu 20.04 system.

My hardware specs are:

Step 1 - Update Nvidia GPU Driver

We will install using the PPA Repository. This will install the latest driver. Run the following command to install the graphics-driver/ppa repository to your system

sudo add-apt-repository ppa:graphics-drivers/ppa

Then install the latest driver. Though we call out 440, it should install the latest (450 in my case)

sudo apt install nvidia-driver-440

You need to reboot to complete the driver installation.

sudo reboot

Check that the driver is Install


Step 2 - Install CUDA Toolkit

Download CUDA Toolkit 11.0

Run the commands provided by Nvidia (also shown below).

sudo mv /etc/apt/preferences.d/cuda-repository-pin-600
sudo dpkg -i cuda-repo-ubuntu2004-11-0-local_11.0.2-450.51.05-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-0-local/
sudo apt update
sudo apt install cuda

If you get an error running the last command to install, you may have conflicting CUDA packages in your /etc/apt/preferences.d directory. Go in and remove the unwanted ones if this occurs.

Step 4 - Install CUDNN 8.0.4

You need an Nvidia account to download cuDNN, so sign up for one (don’t worry, it’s free). Download cuDNN v8.0.4 Library for Linux (x86_64) from the cuDNN Archives. Once, downloaded, we need to unzip the tar file (tar), copy some files to our CUDA Toolkit folders (cp), and change their mode (chmod). To do this, run the following commands:

tar -xzvf cudnn-11.0-linux-x64-v8.0.4.30.tgz
sudo cp cuda/include/cudnn.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn.h /usr/local/cuda/lib64/libcudnn*

Next, we need to update our .bashrc (or .zshrc if you’re using zsh) by adding the path to a few cuda folders to LD_LIBRARY_PATH. Add the following lines to the end of our .bashrc file (or .zshrc file)

export LD_LIBRARY_PATH=/usr/local/cuda-11.0/lib64:$LD_LIBRARY_PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

Step 5 - PIP Install

We will be installing Tensorflow using pip3, which is not installed by default with Ubuntu. You can installed using apt

sudo apt update
sudo apt install python3-pip

We will be using virutal-env, which is also not installed by default with python3 on Ubuntu. If needed, install venv with apt.

sudo apt install python3-venv

We will create a new Virtual Environment with venv and install Tensorflow there. First, let’s make a new virtual environment.

I like to name my virtual environments venv and place them in a directory named after the project name. For example, I create a folder called tensorflow and then create an environment called venv within the tensorflow directory.

mkdir tensorflow
cd tensorflow
python3 -m venv venv

Nex, activate the environment

source venv/bin/activate

Now install tensorflow with pip.

pip install tensorflow

Step 6 - Test It!

We should be good to go at this point. With the environment still activated, start python, import tensorflow, and run the command below to check for GPUs on your system.

import tensorflow as tf

you should see a bunch of output lines and finally a list of your devices. In my case, I get

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

If you get a device listed, you’re good! If you get an empty list, something went wrong. Most likely there is an issue with the CUDNN files not being copied over correctly, or your LD_LIBRARY_PATH is incorrect. that’s it! Enjoy!!