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.
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.
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.
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).
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
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
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
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.
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