Stop Wasting Time: My Journey to Installing TensorFlow with GPU Support

Screenshot of TensorFlow Windows WSL2 installation commands

Have you ever felt like you’re banging your head against a wall trying to get TensorFlow with GPU support to work on WSL2? I have, and after countless hours of frustration, I finally found a solution that works. Here’s my story.

Recently I took up the task of learning ML/AI. Coming from a humble UI/UX background I was enthusiastic to get in this new money making world. A little tutorials on Google and I stumbled upon TensorFlow. I thought lets give it a try.

Disclaimer, I am still not fully clear on what a tensor is. And don’t get me started on other ML terms like confusion matrix. I Google the definition of confusion matrix every morning for the past week

My Struggle with Ubuntu 20.04

If you’ve ever tried to install TensorFlow with GPU support on a WSL, you know how frustrating it can be!

Or maybe it just me who faced these absurd difficulties. But the Nvidia and Github and countless other support tickets at least made me feel I am not alone it this hardship.

Started of with the TensorFlow installation page https://www.tensorflow.org/install. Things seem straight forward enough. Get a Ubuntu distro on your WSL. I choose Ubuntu-20.04 (as it was already something I was working with). Install the supporting softwares like CUDA and cuDNN and off you go. One simple command later you should be ready to embark on you ML journey.

If only that were true.

Post TensorFlow setup I saw one issue. And it just stuck!

tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.

I spent countless hours installing different versions of CUDA and cuDNN, only to find out they didn’t work.

Went through numerous blogs showing how easy it is to install. Poured through multiple Nvidia support forums, to no avail.

Few things I tried doing

• Setup a virtual env with python to install TensorFlow (In hindsight this wasn’t a problem)
• Install latest CUDA and cuDNN versions
• Install the recommended CUDA and cuDNN versions from https://www.tensorflow.org/install/pip#windows-wsl2. Namely CUDA® Toolkit 12.3 and cuDNN SDK 8.9.7.
• Reinstall my Nvidia driver
• Try to run cuda-samples https://github.com/NVIDIA/cuda-samples/tree/master/ to verify my setup. Interesting I found out was there no OpenGL support on WSL yet. So the simple samples work, but the samples like https://github.com/NVIDIA/cuda-samples/tree/master/Samples/2_Concepts_and_Techniques/particles, don’t.

Thing I didn’t try

• Doing a windows native setup as mentioned in https://www.tensorflow.org/install/pip#windows-native. Cause this doesn’t allow me to run the latest TensorFlow version. And which developer in their right mind would settle for any version less than the latest

Caution: TensorFlow 2.10 was the last TensorFlow release that supported GPU on native-Windows. Starting with TensorFlow 2.11, you will need to install TensorFlow in WSL2, or install tensorflow or tensorflow-cpu and, optionally, try the TensorFlow-DirectML-Plugin

• Doing a docker setup.
• And probably n number of other great solutions that I didn’t come accross

The Game-Changer: Ubuntu 24.04

After 2 nights of defeat and what seemed like a 3rd night that won’t end up unfruitful, I decided to gamble on a newer Ubuntu distro, Ubuntu 24.04 to be specific. This decision was a game-changer.

How I Installed TensorFlow on Ubuntu 24.04

Here's a quick guide based on my experience:

• Update Your System:

sudo apt-get update sudo apt-get upgrade

• Install Python’s supporting libraries

sudo apt install python3-pip
pip install --upgrade pip
sudo apt install python3-venv

• Setup a virtual env with python

python3 -m venv tf-gpu
source tf-gpu/bin/activate

• And finally install TensorFlow with GPU support and test

pip install tensorflow[and-cuda]
python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

2025-01-12 06:25:54.846736: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2025-01-12 06:25:55.212793: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1736663155.356183     498 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1736663155.395230     498 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-01-12 06:25:55.770646: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

And that’s it. This one single command did everything for me and gave the output I was trying to get for 3 days.

Conclusion

If you’re struggling to install TensorFlow with GPU support on Ubuntu 20.04, I highly recommend upgrading to Ubuntu 24.04. It saved me a lot of time and frustration. Now, I can finally get on to understanding what actually is a tensor and probably much more.

Happy coding!