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.
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.
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.
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.databatching and prefetching are often as important as the GPU setup.' - Benchmark prediction separately from training because the scaling profile is different.
Related reading
- Keras TimeDistributed Not Masking CNN Model
- Keras TimeDistributed with multiple Inputs in different shapes
- Keras unable to calculate number of parameters in a Keras Custom Layer
- Keras Use the same layer in different models share weights
- Keras TensorFlow Realtime training chart
- Keras TensorFlow Realtime training chart
- Keras Text Preprocessing - Saving Tokenizer object to file for scoring
- Keras Tokenizer num_words doesn't seem to work
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.