Tensorflow with GPU on Windows with WSL and Docker

In this short article we’ll be setting up a desktop environment to run Tensorflow on Windows 10/11 using WSL2 Ubuntu and Docker Desktop for Windows. This article also has a companion video over on YouTube:

If you have trouble completing the steps in the article, please watch the video for a full demonstration of how to complete the steps.

The docker files and sample code used in this tutorial can be found on GitHub at https://github.com/KNuggies/tensorflow-knuggies. If you use the repo, please leave a star :)

In this article, we’ll perform the following steps:

1. Ensure that the NVIDIA Drivers are Installed
2. Install Ubuntu Linux on Windows using WSL2
3. Install Docker Desktop for Windows
4. Run Tensorflow in a Docker Container
5. Improve the Environment with docker-compose

Install the NVIDIA driver for your GPU

This step is probably already complete on your system. If you are not sure, head over to the NVIDIA driver download page and it will walk you through the steps necessary to install the latest drivers. We probably want the game ready driver over the studio driver because most ML machines double as excellent gaming machines.

Install Linux on Windows with WSL2

Before installing WSL2, you need to make sure that you are running Windows 11 or Windows 10 version 2004 and higher (Build 19041 and higher).

You can check your version of Windows 10 by right-clicking the start menu and selecting “System” to display information about your system. Find the Windows specifications as shown below and ensure that you are running build 19041 or higher.

Now that we’ve confirmed that we are running the required version of Windows, let’s install WSL2. Open Windows PowerShell in administrator mode by searching for “PowerShell” in the start menu. Right click the App and select “Run as administrator.”

Once PowerShell opens, run the command:

wsl --install

Follow the on screen instructions. At this point a system reboot will be required. After rebooting, you will create your username and password.

Once the Ubuntu terminal opens, you’ll want to update the system’s software by running the following two commands:

sudo apt update -y
sudo apt upgrade -y

Now that Ubuntu is up to date, let’s go back to Windows to install Docker.

Install Docker Desktop for Windows

To install Docker Desktop, navigate to the site linked above and download Docker Desktop for Windows. It installs similar to any Windows program. Just be sure to look for the “Use WSL 2 instead of Hyper-V (recommended)” option during installation. You want to ensure that this checkbox is selected.

Run Tensorflow in a Docker Container

To run Tensorflow in a Docker Container, we’ll run the following command from the WSL terminal:

docker run -it --rm -p 8888:8888 --gpus all tensorflow/tensorflow:latest-gpu-jupyter

The first time running the command, it will download the Tensorflow container from dockerhub. Once in downloads and runs, click the link to open the Jupyter Notebook server:

http://127.0.0.1:8888/?token=some-gibberish-yours-will-be-different

From here, we can create a new notebook and run the following Python code to confirm that we can see the GPU from the notebook running in the Tensorflow Docker Container.

import tensorflow as tf
tf.config.list_physical_devices()

This is a great start, but if we stop and restart the container, all of our work is gone. Let’s fix that.

Improve the Environment with docker-compose

Now that we have confirmed that the Docker container running Tensorflow has access to the GPU, let’s improve the development environment by creating an easy way to restart our environment and save our work. For this, we’ll need to create a few files. The three files named “requirements.txt”, “Dockerfile” and “docker-compose.yaml” have contents as shown below:

requirements.txt

jupyterlab
pandas
matplotlib

Dockerfile:

FROM tensorflow/tensorflow:latest-gpu
WORKDIR /tf-knugs  # This specifies the directory to work
RUN pip install --upgrade pip
RUN pip install --upgrade -r requirements.txt
EXPOSE 8888
ENTRYPOINT ["jupyter", "lab","--ip=0.0.0.0","--allow-root","--no-browser"]

docker-compose.yaml

services:
  jupyter-lab:
    build: .
    ports:
      - "8888:8888"
    volumes:
      - ./tf-knugs:/tf-knugs
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

Note that the directory “tf-knugs” is used for mapping a volume from the WSL system to the container. These names can be changed to any valid directories. The key is that the WSL directory is on the left side of the “:” and the container directory is on the right side. The directory on the right side should match the “WORKDIR” specified in the Dockerfile.

Once you have these three files saved in the same directory, we can run the new and improved environment by running the following commands to build the container then compose the environment:

docker build .
docker-compose up

Once these steps complete, we’ll see a familiar link to Jupyter Notebooks, but now it opens the newer Jupyter Lab environment that we installed from the “requirements.txt” file.

To stop the environment, go to the WSL terminal used to run “docker-compose up” and press ctrl-C. Alternatively, open a new terminal in the same directory and run:

docker-compose down

Best Learning Resource for Tensorflow and Keras

Now that your environment is set up, keep your progress moving by learning directly from the creator of Keras, François Chollet. Plus, you can help support KNuggies with this affiliate link 🤑:

Conclusion

At this point, we have created a local development environment for Tensorflow with GPU running on Windows WSL Ubuntu using Docker. Congratulations!

If this article has helped you, please follow me here on Medium and cheer this article. Happy coding!