Training machine learning models, especially deep learning models, can be highly resource-intensive. Utilizing GPU acceleration powered by CUDA (Compute Unified Device Architecture) on Ubuntu can significantly reduce the time needed to train large models. In this guide, I’ll walk you through the steps to install CUDA on Ubuntu, with two different installation methods:
- Method 1: Installing CUDA via apt, which simplifies the installation and ensures integration with Ubuntu’s package manager.
- Method 2: Installing CUDA manually from NVIDIA’s official website, which gives you more control over the specific CUDA version you install. After setting up CUDA, we will also demonstrate how to verify that your installation works correctly using popular deep learning frameworks like TensorFlow or PyTorch.
Prerequisites
Before proceeding, ensure that you have the following:
- NVIDIA GPU: This guide assumes you have an NVIDIA GPU that supports CUDA. You can check for compatibility here.
- Ubuntu Version: This guide is based on Ubuntu 20.04 LTS but can be adapted for other versions like 18.04 or 22.04.
- Basic Command Line Knowledge: You will be using terminal commands for most of the steps.
Step 1: Verify Your NVIDIA GPU
First, check whether your system has an NVIDIA GPU and that it's detected by Ubuntu.
Run the following command in the terminal:
lspci | grep -i nvidia
You should see output that lists your NVIDIA GPU. If nothing appears, verify that your system has a compatible NVIDIA GPU installed.
Step 2: Install the NVIDIA Driver
To leverage CUDA, you need to install the appropriate NVIDIA driver. Follow these steps:
Add the NVIDIA graphics drivers PPA repository to your system:
sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update
Find the recommended driver version for your system:
ubuntu-drivers devices
This command will list available drivers. Take note of the recommended driver version.
Install the recommended driver (replace
sudo apt install nvidia-driver-<version>
After the installation, reboot your system:
sudo reboot
Verify the driver installation by running:
nvidia-smi
You should see output listing your GPU details, including the driver version.
Step 3: Install CUDA Toolkit
With the NVIDIA driver installed, you can now proceed to install the CUDA toolkit. Here are two methods to choose from:
Method 1: Install CUDA Using apt
This method simplifies the installation by using Ubuntu's package manager.
Add the NVIDIA CUDA repository and GPG key:
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
sudo sh -c 'echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/cuda.list'
sudo apt update
Install the CUDA toolkit:
sudo apt install cuda
Reboot your system to apply the changes:
sudo reboot
Verify that CUDA is installed by running:
nvcc --version
You should see the CUDA version in the output, confirming the installation.
Method 2: Install CUDA Manually from NVIDIA’s Website
This method gives you more control over the CUDA version you install.
Visit the official CUDA download page.
Select your configuration:
- Operating System: Linux
- Architecture: x86_64
- Distribution: Ubuntu
- Version: 20.04 (or your version)
Download the .deb installer file (choose the network or local installation option).
Install the .deb package by running:
sudo dpkg -i <cuda-repo-package-name>.deb
Replace
- Add the repository key and update the package list:
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
sudo apt update
- Install the CUDA toolkit:
sudo apt install cuda
- After the installation, reboot your system:
sudo reboot
- Add the CUDA binary directory to your system’s PATH by editing your .bashrc file:
echo 'export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc
- Verify the installation:
nvcc --version
Step 4: Test Your CUDA Installation with TensorFlow or PyTorch
Once you have CUDA installed, it's crucial to verify that your machine learning framework can access the GPU. We'll use TensorFlow and PyTorch to ensure CUDA is functioning correctly.
Option 1: Test with TensorFlow
Install TensorFlow with GPU support:
First, make sure you have pip installed, then install TensorFlow:
pip install tensorflow
Test if TensorFlow can detect your GPU by running the following Python code:
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
If CUDA is installed correctly, this should return the number of GPUs available (should be 1 or more).
Option 2: Test with PyTorch
Install PyTorch with GPU support:
Use pip to install PyTorch:
pip install torch torchvision torchaudio
Alternatively, you can visit the official PyTorch installation page for the latest command based on your CUDA version.
Test if PyTorch can use your GPU by running the following Python code:
import torch
print(torch.cuda.is_available())
If CUDA is properly installed, this will return True.
Step 5: Test CUDA with a Sample Program
For further verification, you can also compile and run a sample CUDA program to ensure that CUDA is working properly.
Install the CUDA samples:
sudo apt install nvidia-cuda-toolkit
Compile a sample CUDA program:
Navigate to the CUDA samples directory and build one of the examples:
cd /usr/local/cuda/samples/1_Utilities/deviceQuery
sudo make
Run the sample program:
./deviceQuery
If everything is working correctly, you should see output that details your GPU properties, confirming CUDA is working.
Conclusion
By following this comprehensive guide, you should now have CUDA installed on your Ubuntu system, whether via apt or by manual installation. You’ve also verified the installation using TensorFlow or PyTorch, ensuring that your GPU is ready to accelerate your machine learning models. With CUDA set up, your system is now optimized for efficient deep learning tasks, significantly speeding up training and inference times.
