Skip to content

Instantly share code, notes, and snippets.

@adamelliotfields
Last active July 12, 2024 12:52
Show Gist options
  • Save adamelliotfields/a514e7bdc36215b20a139b811abdae86 to your computer and use it in GitHub Desktop.
Save adamelliotfields/a514e7bdc36215b20a139b811abdae86 to your computer and use it in GitHub Desktop.
Install NVIDIA CUDA and cuDNN on WSL2 for TensorFlow

When neural network frameworks are built, they are dynamically linked to CUDA and cuDNN libraries. These are so or shared object files that are loaded at runtime. The LD_LIBRARY_PATH environment variable tells Ubuntu where to look for these files.

On Windows, these are dll or dynamic link library files.

Why?

GPUs were originally designed for graphics. When you're running a neural network, you're not using the GPU for graphics. CUDA (Compute Unified Device Architecture) is a general-purpose computing on GPUs (GPGPU) platform that allows C-code to run on the GPU. cuDNN (CUDA Deep Neural Network) is a library of primitives like matrix multiplication and convolution that are optimized for GPUs.

To ensure everything works, you want your system to provide the versions of CUDA and cuDNN that your software expects.

Get CUDA and cuDNN versions

Go to TensorFlow's tested build configurations, which is actually a section in the "compile from source" guide. For 2.16.1, it was built with CUDA 12.3 and cuDNN 8.9.

If you already have TensorFlow installed, you can run tf.sysconfig.get_build_info() to see the CUDA and cuDNN versions used at build time.

GPU Driver

In Windows, install the NVIDIA App, which manages your drivers (replacement for GeForce Experience and NVIDIA Control Panel).

There is a support matrix that shows the minimum driver version, but it's only maintained for the latest cuDNN release.

Note that you do not install drivers in Ubuntu (WSL).

Install CUDA Toolkit

Go to CUDA Downloads. Select your configuration, like Linux x86_64 Ubuntu 22.04 deb (network). Note that I'm not using the WSL-Ubuntu distribution.

It will generate a script to download and install a deb package (cuda-keyring):

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-3

Install cuDNN

We need the previous version of cuDNN (8). Go to the cuDNN Archive. The first entry is cuDNN 8.9.7 for CUDA 12.x. We want the Ubuntu 22.04 x86_64 deb version. Note that you have to download it from a browser and be logged into your NVIDIA developer account, so you can't wget it like the CUDA toolkit:

sudo dpkg -i cudnn-local-repo-ubuntu2204-8.9.7.29_1.0-1_amd64.deb
sudo apt update
sudo apt install -y libcudnn8

Update PATH and LD_LIBRARY_PATH

Like the regular PATH, you want to prepend onto LD_LIBRARY_PATH, not overwrite it:

# add CUDA bin to PATH
export CUDA_HOME='/usr/local/cuda'
export PATH="${CUDA_HOME}/bin:${PATH}"

# the `echo...grep` stuff makes it idempotent but isn't required
if [ -d '/usr/lib/wsl/lib' ] && ! echo ":${LD_LIBRARY_PATH}:" | grep -q ':/usr/lib/wsl/lib:' ; then
  LD_LIBRARY_PATH="/usr/lib/wsl/lib:${LD_LIBRARY_PATH}"
fi
if [ -d "$CUDA_HOME/lib64" ] && ! echo ":${LD_LIBRARY_PATH}:" | grep -q ":${CUDA_HOME}/lib64:" ; then
  LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}"
fi
export LD_LIBRARY_PATH="${LD_LIBRARY_PATH%:}" # strip trailing colon

Verify working

First run tf.config.list_physical_devices(). If something isn't right, it will only return the CPU.

Now try a large matrix multiplication operation. It should take a few seconds (depending on cores) on CPU, and be near-instant on GPU:

import tensorflow as tf
size = 10_000
a = tf.random.normal((size, size))
b = tf.random.normal((size, size))
tf.matmul(a, b)

You can also set CUDA_VISIBLE_DEVICES to -1 to force CPU execution.

Uninstall

sudo apt remove -y --purge cuda-keyring cuda-toolkit-12-3 libcudnn8 cudnn-local-repo-ubuntu2204-8.9.7.29
sudo apt autoremove -y

Finally, check /etc/apt/sources.list.d for any leftover files.

To uninstall tensorflow[and-cuda]:

# uninstalls all tensorflow-* and nvidia-* packages in the environment
pip list | grep -E 'tensorflow|nvidia' | awk '{print $1}' | xargs pip uninstall -y
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment