Introduction

Running TensorFlow with GPU support in Docker Compose depends on more than the container image. The host must expose an NVIDIA GPU correctly, Docker must be able to pass that device through, and the Compose file must explicitly reserve GPU access for the service.

If any one of those pieces is missing, TensorFlow usually starts successfully but sees only CPU devices. A good setup verifies GPU visibility at each layer before you try to train a real model.

Check Host Prerequisites First

Before touching Compose, confirm that the host itself can see the GPU. On a Linux host, nvidia-smi should work outside containers.

nvidia-smi

You also need the NVIDIA Container Toolkit so Docker can expose GPU devices to containers. A quick smoke test is running a CUDA image directly:

docker run --rm --gpus all nvidia/cuda:12.9.0-base-ubuntu22.04 nvidia-smi

If this command fails, Compose will fail too. Fix the host driver or container runtime first instead of debugging TensorFlow.

Use a GPU-Capable TensorFlow Image

Your container image must include GPU-enabled TensorFlow and compatible CUDA libraries. A common choice is the official GPU-tagged TensorFlow image.

1services:
2  trainer:
3    image: tensorflow/tensorflow:2.16.1-gpu
4    working_dir: /workspace
5    volumes:
6      - ./:/workspace
7    command: python gpu_check.py
8    deploy:
9      resources:
10        reservations:
11          devices:
12            - driver: nvidia
13              count: 1
14              capabilities: [gpu]

The important part is the devices reservation under deploy.resources.reservations. The capabilities field is required. Use count when you want a number of GPUs, or device_ids when you want specific ones. Do not set both in the same reservation.

Verify TensorFlow Sees the GPU

Create a minimal script and run it through Compose before adding training code.

1import tensorflow as tf
2
3print("TensorFlow version:", tf.__version__)
4print("Built with CUDA:", tf.test.is_built_with_cuda())
5print("Visible GPUs:", tf.config.list_physical_devices("GPU"))
6
7for gpu in tf.config.list_physical_devices("GPU"):
8    tf.config.experimental.set_memory_growth(gpu, True)
9
10with tf.device("/GPU:0"):
11    x = tf.random.normal((1000, 1000))
12    y = tf.matmul(x, x)
13
14print("Matrix shape:", y.shape)

Then start the service:

docker compose up --build

If the output shows at least one physical GPU, your container is wired correctly.

Target a Specific GPU

On multi-GPU machines, pin the service to one device so jobs do not fight over the same hardware.

1services:
2  trainer:
3    image: tensorflow/tensorflow:2.16.1-gpu
4    working_dir: /workspace
5    volumes:
6      - ./:/workspace
7    command: python gpu_check.py
8    deploy:
9      resources:
10        reservations:
11          devices:
12            - driver: nvidia
13              device_ids: ['0']
14              capabilities: [gpu]

You can get device IDs from nvidia-smi. This is useful for shared servers where one Compose project should not consume every GPU on the host.

Keep the Compose File Simple

Once GPU access works, keep the service definition boring. Mount your project, set a working directory, and run one deterministic command. For example:

1services:
2  trainer:
3    image: tensorflow/tensorflow:2.16.1-gpu
4    working_dir: /workspace
5    volumes:
6      - ./:/workspace
7    environment:
8      TF_CPP_MIN_LOG_LEVEL: "2"
9    command: python train.py
10    deploy:
11      resources:
12        reservations:
13          devices:
14            - driver: nvidia
15              count: 1
16              capabilities: [gpu]

The less extra logic you put into startup, the easier it is to tell whether a failure comes from Docker, TensorFlow, or your application code.

Distinguish Compose Problems from TensorFlow Problems

If docker run --gpus all ... nvidia-smi works but TensorFlow still sees no GPU, the issue is usually image compatibility. If even the CUDA container cannot see a GPU, the problem is lower in the stack and unrelated to TensorFlow.

That separation saves time. Too many people start by changing Python packages when the actual issue is that the Docker runtime never exposed the device in the first place.

Common Pitfalls

• Skipping the host-level nvidia-smi check and trying to debug everything from inside TensorFlow.
• Using a CPU-only TensorFlow image and expecting Compose GPU settings to fix it.
• Omitting capabilities: [gpu], which is required in the device reservation.
• Setting both count and device_ids in the same GPU reservation.
• Running large training jobs before confirming with a tiny script that TensorFlow can see the GPU.

Summary

• Verify the NVIDIA driver and container runtime on the host before using Compose.
• Use a TensorFlow image that actually includes GPU support.
• Reserve GPU devices explicitly in the Compose service definition.
• Confirm GPU visibility with a short TensorFlow script before training real models.
• On multi-GPU hosts, pin services to specific devices when isolation matters.