We’re going for Tensorflow 2.4.1 with GPU support on an Ubuntu 20.04 system.
My hardware specs are:
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
nvidia-smi
Download CUDA Toolkit 11.0
Run the commands provided by Nvidia (also shown below).
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget http://developer.download.nvidia.com/compute/cuda/11.0.2/local_installers/cuda-repo-ubuntu2004-11-0-local_11.0.2-450.51.05-1_amd64.deb
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/7fa2af80.pub
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.
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
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
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
tf.config.list_physical_devices('GPU')
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!!