Installing TensorFlow with GPU Acceleration on Linux

Photo courtesy of Coding for Entrepreneurs

So, you want to get started with TensorFlow but don’t know how to set it up. Look no further, you’ve found the right place. This article will explain the steps you should take to end up with a TensorFlow installation with GPU acceleration in a Miniconda environment on Linux. Having a basic understanding of both Linux and Python is helpful, but I’ll be sure tohh explain each of the steps we take. By the end, you should be able to run machine learning models on your local computer at full speed on your GPU. If you don’t have a GPU, you may still find this tutorial helpful but likely won’t get the full benefit out of it.

If you find this tutorial helpful drop a follow.

Table of Contents

1. Warnings
2. Prerequisites
3. Installing Miniconda
4. TensorFlow Environment using Miniconda
5. Installing Nvidia Drivers, CUDA, and cuDNN
6. Installing TensorFlow (finally)
7. Verifying Installation and Comparing CPU vs GPU Speed
8. Troubleshooting
9. TLDR

Obligatory Warning

This tutorial assumes you have a basic level of familiarity with working in the terminal. Please ensure you do not run any commands as sudo or as the root user unless otherwise specified. In various screenshots throughout this tutorial, you may see me as a root user, but that is because I am setting this up in a docker image to avoid touching my current installation of TensorFlow. If that flew over your head — do as I say, not as I do.

This is a lengthy process and because of various hardware differences between each computer, unpredictable errors can occur that may require you to troubleshoot on your own. Additionally, if the configuration isn’t followed in the correct order, you could introduce issues in your system. As such, you should try your best to understand each command and I will do my best to explain each of them. If your system isn’t officially supported, you may also have to change the commands to better fit your computer.

Finally, this is a long, in-depth article, but feel free to jump to the TLDR at the bottom.

Prerequisites

Tensorflow only officially supports Ubuntu 16.04 and newer. However, if you are on another Linux distribution, you should still be able to follow along but it might require more troubleshooting.

We will be doing a lot of work in the terminal, so open that up first (typically Ctrl+Alt+T on Ubuntu).

Once you have it open, first run the following command

cat /etc/os-release

Below is the output when run on my system. The first line, # cat /etc/os-release, is my shell prompt #, followed by the command I run in my shell. The next lines are the output of that command.

You will likely see different values but should overall have a similar output. If you weren’t certain what Linux distribution you were on, you should see its name under NAME. You can also find your version under VERSION_ID.

If you don’t have the os-release file on your system, you probably aren’t on Ubuntu but you can still follow along and hope it works out. You can also check out this link to install TensorFlow in a Docker image, but this may be more complicated depending on your familiarity with Docker.

Now let’s check to ensure you are on a 64-bit system. If you aren’t sure what system architecture you are on, open your terminal again and run:

uname -i

If you don’t see x86_64, you will have limited GPU support but should still be able to run TensorFlow on your CPU. In that case, download the Miniconda version that matches your architecture in the following step.

Finally, TensorFlow only currently supports Nvidia drivers. If you have a different GPU, check out this article.

Installing Miniconda

Let’s start by installing Miniconda.

Miniconda is a lightweight version of Anaconda, a package and environment manager, mainly for Python. In our case, it will help us create a sandboxed environment to install Tensorflow without touching your other Python packages (if you have any). Miniconda is useful because it can typically prevent and resolve dependency conflicts. You can find out more about Miniconda here.

We will download the latest version of Miniconda for Linux for 64-bit systems using the terminal, but you can use the link above if you aren’t on 64-bit.

The command below will download the latest install script for Miniconda and run it. If you are an advanced Linux user, feel free to inspect the script before you run it.

curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -o Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

You will need to press enter to continue, say yes to accept the license terms, and press enter again to install into the default location (typically ~/miniconda3), and yes again to run conda init which will set up your shell to work with conda.

Make sure to close your terminal and reopen it to make sure conda is working. You should see (base) in front of your terminal prompt.

Creating an Environment for TensorFlow using Miniconda

Now we can create our TensorFlow conda environment. Feel free to call your environment something other than tf but keep it short as you may need to type this name often.

The below command will create a Python virtual environment with the name tf and Python version 3.9. You will need to type y and hit enter.

conda create --name tf python=3.9

Note: If you are reading this article far into the future, check on TensorFlow’s website to see what versions of Python they support and change the number accordingly.

Now, you can activate and deactivate your environment with

conda activate tf
conda deactivate

As you can see, after running conda activate tf the words before the prompt changed from (base) to (tf) to indicate I switched from the default Python environment to the one we just installed. Additionally, we can see that when I am in the base environment, the Python version is 3.10, but when I am in our tf environment, Python 3.9 is used, as expected.

When you activate this environment you will have access to all packages installed in it (nothing, at the moment). Once you deactivate, you will return to your default system set of packages.

Installing Nvidia Drivers, CUDA, and cuDNN

TensorFlow requires you to have the latest Nvidia Drivers. You should install them via your system’s package manger. On Ubuntu, you can do this by running sudo apt install nvidia-driver-XXX, replacing XXX with the latest driver version found by running:

apt-cache search --names-only '^nvidia-driver-[0-9]{3}$'

You can also download them from this link, but do so at your own risk as it is always recommended to install drivers from your package manager’s repositories. If you don’t know how to do this, look up how to install the latest proprietary Nvidia drivers on Ubuntu and follow those instructions.

Once you have them, you should be able to run the following command and have a similar output.

nvidia-smi

Now we will install CUDA and cuDNN into our conda environment. The first command activates our tfenvironment, so be sure to replace tfwith whatever you called your environment.

conda activate tf
conda install -c conda-forge cudatoolkit=11.2.2 cudnn=8.1.0

Note: If you are installing this far into the future, check what versions of the cudatoolkit and cudnn you need in the table at this link.

Next, we will create a script that will automatically be run by conda every time you activate your tf environment. This script will just add the path to your newly installed libraries to your environment variables.

conda activate tf

# Create the directories to place our activation and deacivation scripts in
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
mkdir -p $CONDA_PREFIX/etc/conda/deactivate.d

# Add commands to the scripts
printf 'export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\nexport LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${CONDA_PREFIX}/lib/\n' > $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
printf 'export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}\nunset OLD_LD_LIBRARY_PATH\n' > $CONDA_PREFIX/etc/conda/deactivate.d/env_vars.sh

# Run the script once
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

As you can see, before we run those two lines, $LD_LIBRARY_PATH is empty. However, after we run those two lines, a file is created that will automatically be run the next time we activate our TensorFlow environment and will add a value to the LD_LIBRARY_PATH.

Installing TensorFlow (finally)

Now we only have to run two more commands. Make sure you are in the conda environment before running these (tf should be shown in front of your terminal prompt).

pip install --upgrade pip
pip install tensorflow==2.11

Note: Replace 2.11 with the latest version of TensorFlow when you are reading this, which you can find at this PyPI link.

Note: We use pip to install TensorFlow into our conda environment because only pip has the latest updates to TensorFlow, not conda.

Verifying Installation and Comparing CPU vs GPU Speed

Running the following command should display all GPUs that TensorFlow has detected. If you get an error, something along the line failed and you’ll have to debug on your own or check to see if it is covered in the next section.

Note: There is a lot of output TensorFlow logs when you import it, many of which are warnings that we can ignore. Just look for a list on the last line of the output; yours should be similar to [PhysicalDevice(name=’/physical_device:GPU:0', device_type=’GPU’)].

python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

If you want to see the difference between training on the CPU and GPU on TensorFlow, try running the code below. You can also check out this article which shows a similar comparison, in more depth.

Note: Running this will require you to install one more package into your conda environment. It will also autmoatically download a small (38 mb) dataset to ~/tensorflow_datasets.

Run the following commands first to install tensorflow_datasets and then download and run a GitHub Gist I put together to compare the training speeds.

conda activate tf
pip install tensorflow_datasets
curl https://gist.githubusercontent.com/AniAggarwal/052696201e420a873bc404248556cf34/raw/7347c773b485b5d5fe7fadb6ad3f94d693560fe0/speed_comparison.py -O
python speed_comparison.py

As you can see, GPU training took less than half the time (34 seconds vs 84 seconds) of CPU training. This only scales as you start training larger models. So, as you can see, having GPU acceleration is incredibly important.

Troubleshooting

No GPU Found / Empty List Upon Verification

If, when you run the code to verify your installation, only see an empty list on the last line, like so [], then TensorFlow is not detecting your CPU. Try the following steps:

1. Run python -c “import tensorflow as tf; print(tf.test.is_gpu_available())” . If the output is False then your GPU is not being detected.
2. If you are sure you did all the other steps correctly, try installing your graphics driver from this link (same as the one provided in the Nvidia Drivers section).
3. Select your GPU, Linux 64-bit, and Production Branch, then click search and then download. Once the file finishes downloading, you will need to log out of your user account. Be sure to save any work before exiting.
4. Enter a TTY. You can do this by pressing Ctrl+Alt+F3 or any other function key. This will bring you to a CLI asking you to log in.
5. Log in and run killall Xorg. This will stop all graphical applications from running. If you get the error Xorg: no process found then you are probably on Wayland and can skip this step.
6. cdto wherever you downloaded the new graphics driver. Now run chmod +x NVIDIA-Linux-x86_64–X.X.X.run but make sure to replace the X's with the version numbers of your specific file. This will allow the file to be executable.
7. Run the file with ./NVIDIA-Linux-x86_64–X.X.X.run. You may need to run this command with elevated permissions, i.e. sudo ./NVIDIA-Linux-x86_64–X.X.X.run. It will first unpack the file and after a few minutes, it will prompt you a few times. Press ok/yes/allow wherever possible. After, run reboot to restart your computer and try the verify command again like normal.

Couldn’t Find NVVM

If you get an error like the one below (likely because you are on Ubuntu 22.04 like me), run the following commands.

conda activate tf
conda install -c nvidia cuda-nvcc=11.3.58

# Create the directories to place our activation and deacivation scripts in
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
mkdir -p $CONDA_PREFIX/etc/conda/deactivate.d

# Add commands to the scripts
printf 'export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\nexport LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${CONDA_PREFIX}/lib/\nexport XLA_FLAGS=--xla_gpu_cuda_data_dir=${CONDA_PREFIX}/lib/\n' > $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
printf 'export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}\nunset OLD_LD_LIBRARY_PATH\n' > $CONDA_PREFIX/etc/conda/deactivate.d/env_vars.sh

# Run the script once
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

mkdir -p $CONDA_PREFIX/lib/nvvm/libdevice
cp $CONDA_PREFIX/lib/libdevice.10.bc $CONDA_PREFIX/lib/nvvm/libdevice/

• This will add cuda-nvcc to your set of lib files for TensorFlow to use, as well as add them to the script that runs when the tf environment is run. This will override the previous scripts we set up, which is okay.

TLDR

1. Ensure you are on Ubuntu. You can check by running cat /etc/os-release

2. Ensure you are on a 64-bit architecture. You can check by running uname -i

3. Install miniconda. Press enter/type yes as necessary.

curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -o Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

4. Create a new conda env.

conda create - name tf python=3.9

5. Install CUDA and cuDNN

conda activate tf
conda install -c conda-forge cudatoolkit=11.2.2 cudnn=8.1.0

6. Ensure you have latest Nvidia drivers. Check using nvidia-smi. When you run the verification command later on, if don’t see your GPU, you will need to revisit this step. Check the Troubleshooting section if that is the case.

7. Create a script that is activated every time your conda env is activated/deactivated to add CUDA libraries to your path.

conda activate tf

# Create the directories to place our activation and deacivation scripts in
mkdir -p $CONDA_PREFIX/etc/conda/activate.d
mkdir -p $CONDA_PREFIX/etc/conda/deactivate.d

# Add commands to the scripts
printf 'export OLD_LD_LIBRARY_PATH=${LD_LIBRARY_PATH}\nexport LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${CONDA_PREFIX}/lib/\n' > $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
printf 'export LD_LIBRARY_PATH=${OLD_LD_LIBRARY_PATH}\nunset OLD_LD_LIBRARY_PATH\n' > $CONDA_PREFIX/etc/conda/deactivate.d/env_vars.sh
# Run the script once
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

8. Install TensorFlow in your conda env.

conda activate tf
pip install --upgrade pip
pip install tensorflow==2.11

9. Verify it works. Look for a list of GPU devices.

python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

10. If you want to compare the difference between running on CPU and GPU, run the following commands:

conda activate tf
pip install tensorflow_datasets
curl https://gist.githubusercontent.com/AniAggarwal/052696201e420a873bc404248556cf34/raw/7347c773b485b5d5fe7fadb6ad3f94d693560fe0/speed_comparison.py -O
python speed_comparison.py

If you run into any issues, check the Troubleshooting section or go back and read skim the more detailed article.