Simple guide for installing TensorFlow (GPU version) on WSL2 | by Momchil Georgiev | Jan, 2024

Finito!

This article is directed to people who are fairly new to TensorFlow, as one can probably guess. It’s also assumed that the reader has limited experience with WSL2 and Linux, so don’t panic pre-emptively!

The main reason for writing this article is that I find the whole process of setting up TensorFlow (GPU enabled version) needlessly complicated. So this article is intended to be simple and (I hope) easy to follow. For more comprehensive dive into the different subjects mentioned here I think there is enough information on the net including the free access to ChatGPT and other LLMs like Bard.

I will now walk you through the whole process of setting-up TensorFlow step by step and explain the main parts to the best of my abilities.

Here is a walkthrough on what we will cover:

1. Installing WSL2 in which we will configure TensorFlow
2. Install Miniconda and create virtual environment for TensorFlow
3. Download and install Tensorflow
4. Discuss and troubleshoot

First, make sure to have installed NVIDIA Game ready drivers for your GPU on your WINDOWS machine. Just type “NVIDIA Game ready drivers” in Google and download them. Unfortunately AMD GPUs are not supported since the libraries TensoFlow uses are by… NVIDIA. Sorry guys.

Second thing which is really important — DO NOT FOR GOD SAKE INSTALL CUDA AND CUDA TOOLKIT ON YOUR WINDOWS SYSTEM!!!

Cuda and cuda toolkit are ALREADY downloaded in your Ubuntu setup (read 3. in the link provided) from the start! Downloading those dependencies on your PC will cause a lot of trouble (speaking from experience)! Just install the Game ready drivers and continue with the tutorial.

Setting up WSL2 — Windows Subsystem for Linux is required do use the newest versions of TF which has GPU support. WSL is simply put a virtualization of the Linux kernel, allowing you to run Linux commands and programs from your Windows PC. The simplest way is to get it from the MS Store or through PowerShell. After installing it your PC has to reboot.

You will use WSL to install Ubuntu in a virtual environment. Ubuntu is officially supported both by NVIDIA and Microsoft and going with it will give you less headaches. This means that hardware drivers work pretty much out of the box and most of the software part is already done, although this doesn’t mean that other distros won’t run TF.

After installing WSL it is easy to install Ubuntu. Run this in your cmd:

wsl install -d Ubuntu

This installs Ubuntu. Reboot after installing it.

Congrats, now you have virtualized Linux distro on your PC. Now to open and configure it to run TensorFlow type wsl in cmd:

wsl

Since Ubuntu is your only distro, it will start automatically.

Now you are in Ubuntu’s terminal. Our goal is to install TensorFlow in a separate environment since it uses older versions of dependencies and this way it is easier to manage and change things in your system.

Check your gpu setup by typing this in the Ubuntu terminal:

nvidia-smi

Should look something like this:

If this shows up, it means that you also have CUDA and CUDA Toolkit installed. If it doesn’t, you probably don’t have GPU drivers on your Windows PC. Install them, reinstall WSL (just in case) and Ubuntu and continue with the tutorial.

TensorFlow officially says to use Miniconda. To download it run this command which downloads and installs it from the terminal:

curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -o Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh

To create an environment run:

conda create --name tf-wsl python=3.9

This creates an environment called tf-wsl in which we will install TF. To access it type this into Ubuntu’s terminal:

conda activate tf-wsl

If it works you should have (tf-wsl) in your terminal. Sort of like this:

Next we will update pip and finally… download TensorFlow! To do that type in Ubuntu terminal this:

pip install --upgrade pip
pip install tensorflow[and-cuda]

To validate everything works call this command which simply opens python and runs TF code:

python3 -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

If the output has this at the end, you are done!

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

This is the simple part! Open VS Code, press the bottom left button (the one which looks like “><”) and select “Connect to WSL”.

After clicking it you will have access to the files in your Ubuntu distro. To access files and actually do your work you have to know how the file structure works in WSL. The directory of the distro can be found in File Explorer under “Linux” on the sidebar:

Clicking “Ubuntu” opens up the main directory of the distro:

I suggest doing your work in ./home since you can create and move files in it with no need for setting up permissions.

You can copy or create a notebook in this directory to test it by going to “File > Open Folder…” in VS Code and opening the directory you want.

Here is an example with me. I have a working directory “dl” which has my work in it:

Opening it will put you in your familiar spot!

Try loading a library. This will make VS Code to recommend you download some other packages and extensions related to jupyter notebook and python that you will need. Don’t worry! It has happens automatically and just click “install” where it tells you to.

I now hope I will save you some time and anxiety.

1. PATH Variables

The most common problem is when you have everything done and close VS Code. After turning it on and loading into WSL again for some reason you cannot find your environment or you can but TF is not installed. This is caused by the PATH variables not being set up properly. See how to set them up (this goes outside the scope of this article). ChatGPT, Stack Overflow and Google can help a lot in here.

2. You don’t see your GPU when validating you have it (step 7 from the list at the beginning)

This means you haven’t installed your GPU drivers right. Again, install NVIDIA Game ready drivers, then WSL and Ubuntu and everything should be there working.

3. You were training a model. You ran out of memory and decided to kill VS Code with Task manager and now WSL constantly disconnects you.

This is done by some weird caching thing. The only way (to my knowledge) to resolve this is unfortunately to reinstall Ubuntu or WSL in general.

Hope this helped you. If there is anything wrong or has to be added please let me know. I would love some feedback!