Keras
GPU
Deep Learning
Neural Networks
Machine Learning

How to use Keras with GPU?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Keras does not usually need special GPU-only model code. If TensorFlow can see a compatible GPU, the expensive numeric parts of your model will generally run there automatically, so most setup problems are really environment and verification problems.

Verify That TensorFlow Sees The GPU

Before changing any model code, confirm that the runtime can actually detect a GPU:

python
1import tensorflow as tf
2
3print(tf.__version__)
4print(tf.config.list_physical_devices("GPU"))

If the list is empty, Keras will fall back to the CPU. That does not mean your model definition is wrong. It usually means the machine runtime, drivers, or TensorFlow installation are not aligned for GPU execution.

When a GPU is present, enabling memory growth is often a good first step:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4for gpu in gpus:
5    tf.config.experimental.set_memory_growth(gpu, True)
6
7print(gpus)

Do this before creating tensors or models. Otherwise TensorFlow may initialize the device too early and reject the configuration change.

Standard Keras Training Usually Just Works

Once TensorFlow sees the device, ordinary Keras code usually needs no special switch:

python
1import numpy as np
2import tensorflow as tf
3from tensorflow import keras
4from tensorflow.keras import layers
5
6x_train = np.random.rand(2048, 20).astype("float32")
7y_train = np.random.randint(0, 2, size=(2048, 1)).astype("float32")
8
9model = keras.Sequential(
10    [
11        layers.Input(shape=(20,)),
12        layers.Dense(64, activation="relu"),
13        layers.Dense(32, activation="relu"),
14        layers.Dense(1, activation="sigmoid"),
15    ]
16)
17
18model.compile(
19    optimizer="adam",
20    loss="binary_crossentropy",
21    metrics=["accuracy"],
22)
23
24model.fit(x_train, y_train, epochs=3, batch_size=64)

If the environment is correct, the heavy tensor operations used during training will be placed on the GPU automatically.

Confirm Device Placement While Debugging

If you want proof that operations are actually landing on the GPU, ask TensorFlow to log placement:

python
import tensorflow as tf

tf.debugging.set_log_device_placement(True)

That setting can be noisy, but it is helpful when you suspect that expensive operations are silently running on the CPU.

Another practical check is timing. If a large model runs no faster than CPU-only training, the bottleneck may be in data loading or preprocessing rather than the model step itself.

Use One GPU Or Many GPUs Deliberately

On machines with multiple GPUs, you may want to limit visibility to one device:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4if gpus:
5    tf.config.set_visible_devices(gpus[0], "GPU")
6    print(tf.config.get_visible_devices("GPU"))

For multi-GPU training, do not manually split the batches yourself. Use a distribution strategy:

python
1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5strategy = tf.distribute.MirroredStrategy()
6
7with strategy.scope():
8    model = keras.Sequential(
9        [
10            layers.Input(shape=(20,)),
11            layers.Dense(64, activation="relu"),
12            layers.Dense(1, activation="sigmoid"),
13        ]
14    )
15    model.compile(optimizer="adam", loss="binary_crossentropy")

This keeps the model-building code straightforward while TensorFlow coordinates replicas behind the scenes.

Tune The Workload, Not Just The Hardware

A visible GPU is only the starting point. Real training speed also depends on batch size, input pipeline performance, preprocessing, and whether the model is large enough to benefit from accelerator hardware.

Small toy models can run only marginally faster on a GPU because kernel launch overhead and host-to-device transfers may dominate. If training feels slow, profile the full pipeline instead of assuming the GPU itself is the problem.

Common Pitfalls

  • Assuming Keras needs a separate GPU-specific training API.
  • Verifying package installation but never checking tf.config.list_physical_devices("GPU").
  • Setting memory growth after TensorFlow has already touched the device.
  • Expecting tiny models or tiny datasets to show dramatic GPU speedups.
  • Ignoring data loading bottlenecks when the GPU looks underutilized.

Summary

  • Keras usually uses the GPU automatically through TensorFlow.
  • Verify device visibility first before changing model code.
  • Standard model.fit(...) code is normally enough once the environment is correct.
  • Use device placement logs, memory growth, and distribution strategies when you need more control.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.