I have been setting up multiple Python machine learning environments based on NVIDIA GeForce GPUs. Starting with a fresh Ubuntu installation, you will need to go through the installation of multiple NVIDIA packages to get your machine ready to communicate with the GPU(s) properly. The steps here are an overview of the installation process.

Specifications

Here are my system specifications:

  • Ubuntu 22.04 LTS
  • NVIDIA GeForce RTX 3090 x2

Installation

Since drivers and toolkits follow different versioning schemes, I have added the latest versions as well just to make it easier to know what is what!

1. NVIDIA Drivers

It is recommended to install the drivers via the OS package manager. You can follow the steps here.

2. CUDA Toolkit

You can find the official CUDA Toolkit installation instructions on Linux here. One-time skimming through the entire documentation is recommended as there are a variety of dependencies and settings depending on your OS and the expected use case. In a high-level view, the process at this stage can be broken down into three main steps: 1. pre-installation actions, 2. main installation, and 3. post-installation actions

2.1 Pre-Installation Actions

As long as you are having a typical setup, there is no special step to take here.

2.2 Package Manager Installation

I prefer the Network Repo (over the Local Repo) installation as it gives you the ability to update through the APT package manager in the future. Follow the instructions for Ubuntu here. In summary, you need to install the cuda-keyring and then the CUDA SDK. Once you finish, you need to reboot your system.

2.3 Post-Installation Actions

Some finishing touches can be found here. Make sure you have updated the PATH variable. Adding the export shell command to your shell configuration (e.g. bashrc or zshrc) file is recommended to keep the PATH variable consistent throughout a shell session.

3. cuDNN

The steps here are similar to the ones from the previous step. Install the package according to the documentation here.

Finale

Once you have all the packages installed, you can verify your installation. In the case of machine learning you can follow the quick checks with the PyTorch and TensorFlow libraries from my other post.