TensorFlow GPU on Ubuntu 20.04

3 min read 20-10-2024
TensorFlow GPU on Ubuntu 20.04

If you're looking to harness the power of TensorFlow with GPU acceleration on Ubuntu 20.04, you've come to the right place. This guide will take you through the installation process step-by-step, ensuring you set up your environment for optimal performance.

Why Use TensorFlow with GPU?

TensorFlow is a powerful open-source library for machine learning and artificial intelligence, and leveraging GPU (Graphics Processing Unit) capabilities can dramatically speed up computations. With TensorFlow running on GPU, you can train models faster and handle larger datasets more efficiently, which is vital for deep learning applications.

Prerequisites

Before diving into the installation, ensure you have the following:

  • An Ubuntu 20.04 machine.
  • NVIDIA GPU that supports CUDA (Compute Unified Device Architecture).
  • Basic knowledge of using the terminal.

Step 1: Update Your System

First, open your terminal and update your package list to ensure all your system packages are current:

sudo apt update && sudo apt upgrade -y

Step 2: Install NVIDIA Driver

Install the NVIDIA driver that matches your GPU. You can do this by running:

sudo ubuntu-drivers autoinstall

After installation, reboot your system:

sudo reboot

You can verify the installation of the NVIDIA driver using:

nvidia-smi

This command will display information about your GPU, including the driver version and GPU utilization.

Step 3: Install CUDA Toolkit

Next, download the CUDA toolkit. You can find the latest version compatible with TensorFlow on the NVIDIA CUDA Toolkit page. Choose the correct version for Ubuntu 20.04.

Once downloaded, install the CUDA toolkit with:

sudo sh cuda_<version>_linux.run

Follow the on-screen instructions, and make sure to set the environment variables. Add the following lines to your .bashrc file:

echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

Step 4: Install cuDNN

cuDNN (CUDA Deep Neural Network library) is essential for deep learning. You can download the appropriate version of cuDNN from the NVIDIA cuDNN page. Make sure to choose the version compatible with the version of CUDA you installed.

After downloading, install cuDNN using:

tar -xzvf cudnn-*.tgz
sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod 777 /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

Step 5: Install TensorFlow with GPU Support

With CUDA and cuDNN installed, the next step is to install TensorFlow with GPU support. You can do this via pip. It's recommended to create a virtual environment:

python3 -m venv tf_gpu
source tf_gpu/bin/activate
pip install --upgrade pip
pip install tensorflow

To verify the installation, you can run the following command in Python:

import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

If correctly installed, you should see the number of GPUs available for TensorFlow.

Conclusion

You are now ready to harness the full power of TensorFlow with GPU support on your Ubuntu 20.04 machine. This setup will significantly accelerate your machine learning projects. Remember that keeping your software updated and following best practices for GPU computing will help you achieve the best results.

Additional Resources

With this guide, you can now effectively set up and utilize TensorFlow with GPU on your Ubuntu system. Happy coding!