Installation of Tensorflow2 with GPU support is easy and the only complication can be arisen from the CUDA compability which in turns depends on the Nvidia driver version. Before going farther, please check if your Nvidia Video Card is compatible with the required versions that are defined in this gist, use this link.
Tensorflow offers in its website a table of the compatibility between libraries for the target OS. You can visit that website in the following COMPATIBILITY TABLE that points to the Tensorflow installation from source for linux. For the time writing this gist, Tensorflow2 v2.17.0 requires CUDA v12.3 and CUDNN v8.9. As CUDA v12.6 gives support for Ubuntu 24.04, I've replaced CUDA v12.3 with CUDA v12.6, but maintaining CUDNN v8.9. It is really important to match the exact version of CUDNN, otherwise tensorflow will have problems loading the shared libraries as not finding the correct version.
CUDA version also requires for a minimum Nvidia driver version. In our case, CUDA v12.6 requires Nvidia driver to be superior to 525.60.13 and you can check that information from CUDA website that can be found in this CUDA-COMPATIBILITY TABLE.
For the CUDA installation we must visit Cuda website and select the exact version. Nowadays, it is even not the latest one so you have to navigate through the legacy releases to find the exact version. If you do not find the correct link, google "cuda archive" to find the website with previous releases. There are several ways to install CUDA, I prefer the runfile(local) option. So, from the form you must choose Linux > x86_64 > Ubuntu > 24.04 > runfile(local). There is a text box below showing the commands to follow for the installation. In our case, we have the following instructions:
$ wget https://developer.download.nvidia.com/compute/cuda/12.6.1/local_installers/cuda_12.6.1_560.35.03_linux.run
$ sudo sh cuda_12.6.1_560.35.03_linux.run
Possible Errors: If the installation fails with the following message ´Installation failed. See log at /var/log/cuda-installer.log for details´, it means that your system is using current nvidia-drivers and this operation cannot be permitted. Remember, that this installation is not only installing the CUDA Toolkit, but also the nvidia drivers. We could manage this error by installing only the toolkit and it can be done using the following arguments in the previous command:
$ sudo sh cuda_12.6.1_560.35.03_linux.run --toolkit --silent --override
For the CUDNN installation, we have to start downloading the CUDNN version from Nvidia (an account is required to be logged in the website to proceed with the dowload) through this link. In the website look for the CUDNN version compatible with your CUDA version.
Next steps are really well explained by LearnOpenCV blog in the following post Installing Deep Learning Frameworks on Ubuntu with CUDA support. Basically, we have to uncompress the file and copy the extracted libs to the cuda installation folder.
$ tar -xvf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz
$ sudo cp -P cudnn-linux-x86_64-8.9.7.29_cuda12-archive/lib/* /usr/local/cuda/lib64/
$ sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/include/* /usr/local/cuda/include/
Now, we have to configure some enviroment variables to help the system locate these libraries:
$ echo 'export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"' >> ~/.bashrc
$ echo 'export CUDA_HOME=/usr/local/cuda' >> ~/.bashrc
$ echo 'export PATH="/usr/local/cuda/bin:$PATH"' >> ~/.bashrc
Reload the file:
$ source ~/.bashrc
First of all, it is always recommended to use virtual environments. There is a major update in Ubuntu 24.04 with python 3.12, so it is recommended to use pipx:
Install PIPX:
-------------
$ sudo apt install pipx
$ pipx ensurepath
Install Virtualenv and Wrapper:
-------------------------------
$ pipx install virtualenv
$ pipx install virtualenvwrapper
Configure ~/.bashrc:
--------------------
$ nano ~/.bashrc
... #Append these lines, but changing the user folder name $USER
#VIRTUAL ENVIRONMENT
export WORKON_HOME=~/.virtualenvs
export VIRTUALENVWRAPPER_PYTHON=/home/$USER/.local/share/pipx/venvs/virtualenvwrapper/bin/python3
source /home/$USER/.local/share/pipx/venvs/virtualenvwrapper/bin/virtualenvwrapper_lazy.sh
...
$ source ~/.bashrc
Create virtualenv and activate:
-------------------------------
mkvirtualenv tensorflow
workon tensorflow
Once created the virtual environment, activate it and install Tensorflow using PyPy:
$ pip3 install tensorflow //By default it will install the latest estable version of Tensorflow
To check the correct installation import the tensorflow package and run a sequence to test if the GPU is available:
$ python3 //To enter in the python interpreter
> import tensorflow as tf
> tf.test.is_gpu_available()
If you sucess you will receive a ´True´ response.
import tensorflow as tf
from tensorflow.keras import datasets, layers, models
# Check TensorFlow version
print(f"TensorFlow version: {tf.__version__}")
# Print CUDA and cuDNN versions
print(
"Num GPUs Available: ",
len(tf.config.experimental.list_physical_devices("GPU")),
)
print("CUDA Version: ", tf.test.is_built_with_cuda())
print("GPU Device Name: ", tf.test.gpu_device_name())
# Enable mixed precision training
# This can speed up training on newer GPUs
policy = tf.keras.mixed_precision.Policy("mixed_float16")
tf.keras.mixed_precision.set_global_policy(policy)
# Load and preprocess the MNIST dataset
mnist = datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = (
x_train.astype("float32") / 255.0,
x_test.astype("float32") / 255.0,
)
# Reshape the data to add the channel dimension
x_train = x_train.reshape(x_train.shape[0], 28, 28, 1)
x_test = x_test.reshape(x_test.shape[0], 28, 28, 1)
# Create the model
model = models.Sequential(
[
layers.Conv2D(32, (3, 3), activation="relu", input_shape=(28, 28, 1)),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation="relu"),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation="relu"),
layers.Flatten(),
layers.Dense(64, activation="relu"),
layers.Dense(10, activation="softmax"),
]
)
# Compile the model
model.compile(
optimizer="adam",
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
# Train the model
history = model.fit(
x_train, y_train, epochs=5, validation_split=0.2, batch_size=128
)
# Evaluate the model
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print(f"\nTest accuracy: {test_acc:.3f}")
Download the .deb file and run the following installation instructions offered by CUDA:
$ sudo dpkg -i cuda-repo-ubuntu1804_10.0.130-1_amd64.deb
$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
$ sudo apt-get update
$ sudo apt-get install cuda-toolkit-10-0 //Be aware!! I have change this line to only install the respective toolkit.
Important, before installing cuda package make sure that this package is not upgrading your Nvidia drivers, otherwise it would mess up the installation. In fact, we only want to install the CUDA Toolkit not the CUDA package that contains the toolkit and the Nvidia drivers. If you want to do not fail this step launch the proposed sentence: ´sudo apt-get install cuda-toolkit-10-0´
Possible Errors: If you have problems installing cuda-repo-ubuntu...deb because you have previously installed versions of cuda, you can find the previous packages using ´dpkg -l | grep cuda*´ and delete them with ´dpkg -r package_name´.