Use A GPU With TensorFlow on Windows

A cow sat at an office desk with a big, red stapler on the table

Adie Does Typical Adie Things

I started a simple project with just one, simple intention - do the TensorFlow equivalent “Hello, world”.

How hard could it be?

Erm, turns out it wasn’t hard at all, as it happens. But with me being me, I found a small issue that bugged me for a super long time - it wasn’t running on my GPU.

Windows 11 Causes a Headache

I use Windows on main (I am a gamer, so that’s my excuse). After what can only be described as “much faffing about” I found out that anything above version 2.10 of TensorFlow is not supported. For no real, logical reason, I was against using an old version and decided to find an alternative solution.

Funnily enough, both solutions to get it running in Windows involved sticking some Linux in it. One involved using WSL 2 (Windows Subsystem For Linux). And then along came Docker.

Let’s Go Whaling 🐳

The Docker platform allows you to run ‘containers’ which encapsulate everything needed to run pieces of software - that’s the code, runtime, libraries, and dependencies. It’s kinda like a lightweight, portable virtual environment. Many companies have official Docker images available for use. You build an image into a container and can start and stop the containers whenever you want.

TensorFlow has a Docker image which handily comes with a pre-configured Linux environment also containing Python and all the dependencies needed. The Nvidia libraries for GPU-accelerated deep learning, CUDA, and cuDNN, are also contained in the package.

What happens with your Python scripts is that they are copied into the Docker image, and when you run a Docker container from the image, your script runs in this container. If you have other Python packages that your script depends on (dependencies), you can add them to a requirements.txt file.

Prepare to Build the TensorFlow Image

We’ll need a Dockerfile and docker-compose.yml file. Use this in the former:

FROM tensorflow/tensorflow:latest-gpu

ENV PYTHONDONTWRITEBYTECODE 1

ENV PYTHONUNBUFFERED 1

WORKDIR /app

ADD . /app

RUN pip install --no-cache-dir -r requirements.txt

This does the following:

Specifies the parent image
Stops Python from generating .pyc files (they aren’t needed here and can cause problems)
Turns off buffering for easier container logging, debugging
Sets the working directory in the container to /app
Adds the current directory contents into the container at /app
Installs additional packages needed. Docker uses an internal cache when building the image.

This docker-compose.yml file does a few things:

version: "3.8"
services:
  app:
    build:
      context: .
      dockerfile: Dockerfile
    volumes:
      - .:/app
    command: /bin/bash -c "tail -f /dev/null"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

Specifies that we’re going to use version 3.8 of Docker Compose.
The services section defines a service named app built using the specified Dockerfile in the current directory (.).
volumes mounts the current directory onto the /app directory within the container.
command contains a line that will keep the container running, useful for debugging and other interactions.
deploy specifies deployment-related configurations. In this case it reserves all available GPU devices using the Nvidia GPU driver. (What, you only have ONE graphics card?)

Build and Run It

You will need Docker Desktop, so go and download it and open it.

In your normal terminal, build the image and run the container with the command:

docker-compose up --build

Make sure you do that in the same directory you created the two, Dockerfile and docker-compose.yml.

You can shut down the container with docker-compose down and start it again without rebuilding it with docker-compose up. Docker Desktop also has some icons for you to click on if you prefer using a GUI.

So you’ve done all that and you’ve got some TensorFlow Python code ready to run. But how do you do that inside the container? You attach a terminal to it by running this command:

docker-compose exec app bash

After doing that you will be in a brand new terminal inside the container, like magic, and TensorFlow should now utilise your flashy GPU instead of your CPU. 🤗