machine learning
deep learning
keras
tensorflow
multi-gpu

Keras Tensorflow Prediction on multiple gpus

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

Yes, Keras running on TensorFlow can use multiple GPUs for prediction, not just training. The usual mechanism is tf.distribute.MirroredStrategy, which replicates the model across local GPUs and splits batches automatically.

The important detail is that multi-GPU inference only helps when the model and batch size are large enough to amortize coordination overhead. Small models or tiny batches often run just as fast on one GPU.

How Multi-GPU Prediction Works

TensorFlow uses data parallelism. Each replica receives part of the input batch, computes predictions locally, and TensorFlow combines the outputs in the original order.

For prediction workloads, this means:

  • the model is copied to each GPU
  • the input batch is divided across replicas
  • outputs are gathered back into one result tensor

This is convenient because the same Keras model can usually be trained and served with the same distribution strategy on one machine.

A Runnable Keras Example

The following example builds a simple model, trains briefly on synthetic data, and then performs prediction through a mirrored strategy. It also prints the visible GPUs so you can confirm the environment.

python
1import numpy as np
2import tensorflow as tf
3
4gpus = tf.config.list_physical_devices("GPU")
5print("GPUs:", gpus)
6
7strategy = tf.distribute.MirroredStrategy()
8print("Replicas:", strategy.num_replicas_in_sync)
9
10with strategy.scope():
11    model = tf.keras.Sequential([
12        tf.keras.layers.Input(shape=(10,)),
13        tf.keras.layers.Dense(32, activation="relu"),
14        tf.keras.layers.Dense(1, activation="sigmoid"),
15    ])
16    model.compile(optimizer="adam", loss="binary_crossentropy")
17
18x_train = np.random.rand(4096, 10).astype("float32")
19y_train = np.random.randint(0, 2, size=(4096, 1)).astype("float32")
20
21model.fit(x_train, y_train, epochs=2, batch_size=256, verbose=0)
22
23x_test = np.random.rand(2048, 10).astype("float32")
24predictions = model.predict(x_test, batch_size=512, verbose=0)
25
26print(predictions.shape)

If the machine has only one GPU, the code still works. MirroredStrategy simply uses the devices it can find.

When It Helps

Prediction scales best when all of the following are true:

  • the model is large enough to keep each GPU busy
  • the batch size is large enough to split effectively
  • preprocessing is not the bottleneck
  • data transfer from host memory to GPU is not dominating the runtime

This is why image classification and transformer inference often benefit from multi-GPU prediction, while very small tabular models often do not.

If your pipeline feeds one example at a time, GPU utilization will stay low no matter how many devices are present. The fix is usually to batch requests and use tf.data so the input pipeline keeps the devices busy.

Input Pipelines Matter

A slow Python loop can erase any advantage from multiple GPUs. Prefer tf.data.Dataset so TensorFlow can prefetch and batch efficiently.

python
1dataset = tf.data.Dataset.from_tensor_slices(x_test)
2dataset = dataset.batch(512).prefetch(tf.data.AUTOTUNE)
3
4predictions = model.predict(dataset, verbose=0)

That small change often matters more than adding another GPU. Inference throughput is a system problem, not just a model problem.

Training and Prediction Are Not Identical

People sometimes assume that because multi-GPU training works, multi-GPU prediction will show the same speedup. That is not guaranteed.

Training has expensive forward and backward passes, so distributing the work often pays off. Prediction only has the forward pass, which means:

  • there is less computation to amortize replication overhead
  • host-to-device transfer can dominate
  • output gathering becomes more visible in profiles

You should therefore benchmark prediction separately rather than reusing training assumptions.

Common Pitfalls

  • Expecting speedup from multiple GPUs with a tiny batch size.
  • Forgetting that the input pipeline can bottleneck inference before the GPUs do.
  • Creating the model outside strategy.scope() when the intent is distributed execution.
  • Measuring one request at a time instead of realistic batched workloads.
  • Assuming replica count alone determines performance. Memory bandwidth, preprocessing, and model shape matter too.

Summary

  • Keras with TensorFlow can run prediction on multiple GPUs through tf.distribute.MirroredStrategy.
  • Multi-GPU inference uses data parallelism and batch splitting.
  • The biggest gains appear on large models and sufficiently large batches.
  • 'tf.data batching and prefetching are often as important as the GPU setup.'
  • Benchmark prediction separately from training because the scaling profile is different.

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.