How to speed up Tensorflow 2 keras model for inference?
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
Inference speed problems in TensorFlow 2 usually come from a small number of causes: too much Python overhead, the wrong batch size, or a model format that does not match the deployment target. The fastest path is to measure first, then remove overhead in the serving path before reaching for heavier optimizations.
Benchmark the Real Inference Path
Before changing code, benchmark the exact call pattern your application uses. A surprising number of slow deployments are measuring model.predict() in a notebook while production uses single-request calls from a web server.
Always warm up first. The first few calls often include tracing and one-time setup cost that should not be mixed into steady-state latency.
Reduce Python Overhead
For low-latency serving, calling the model directly is often faster and simpler than using model.predict(). Wrapping inference in tf.function lets TensorFlow stage more work in the graph and reduces Python overhead per request.
This is especially helpful when inference is called many times with the same tensor shapes. If your requests have wildly different shapes, retracing can eat the performance win, so try to keep shapes stable where possible.
Batch Requests When Throughput Matters
Latency and throughput are different goals. If you want maximum requests per second, batching is usually the biggest win because it makes better use of vectorized kernels and GPU hardware.
There is no universal best batch size. Small batches can be better for interactive latency, while moderate or large batches often win on throughput. Measure on the hardware you actually deploy.
Use an Optimized Deployment Format
If the model will run on mobile, edge, or CPU-bound production systems, converting it can produce a bigger gain than tweaking the Python call site.
TensorFlow Lite is a common option:
This step can reduce model size and enable backend-specific optimizations. On supported accelerators, other runtimes such as TensorRT may give even better results. The correct choice depends on whether you deploy to CPU, GPU, mobile, or edge hardware.
Let TensorFlow Optimize the Graph
TensorFlow can apply graph-level optimizations automatically, and enabling XLA is sometimes worth testing for stable workloads.
XLA is not a guaranteed win for every model, but it is one of the easiest experiments to run. Keep it behind a benchmark and only keep it if it improves your real workload.
Mixed precision can also help on supported GPUs:
Use that carefully and re-check numerical behavior, especially if post-processing assumes float32.
Simplify the Serving Path
A fast model can still feel slow if the surrounding pipeline wastes time converting data. Common examples include:
- decoding images repeatedly in Python,
- converting between NumPy arrays and tensors for every request,
- resizing inputs one image at a time,
- loading the model inside the request handler.
Load the model once at startup, preallocate where possible, and keep preprocessing close to TensorFlow ops when it is practical.
Common Pitfalls
- Timing the very first inference call and mistaking startup cost for normal latency.
- Using
model.predict()in a tight serving loop where directmodel(batch, training=False)would be lighter. - Assuming GPU is always faster, even for tiny single-item batches where CPU can win.
- Converting to TensorFlow Lite or enabling XLA without benchmarking the actual deployment target.
- Ignoring preprocessing overhead, which can dominate total response time even when the model itself is fast.
Summary
- Benchmark the exact inference path you plan to serve, including warmup.
- Reduce Python overhead with direct model calls and
tf.function. - Tune batch size based on whether you care more about latency or throughput.
- Consider optimized formats such as TensorFlow Lite for deployment-specific gains.
- Measure end-to-end time, not just raw model execution, because preprocessing often becomes the real bottleneck.
Related reading
- How to speedup rnn training speed of tensorflow?
- How to Split the Input into different channels in Keras
- How to stack multiple lstm in keras?
- How to tell if tensorflow is using gpu acceleration from inside python shell?
- How to store best models checkpoints, not only newest 5, in Tensorflow Object Detection API?
- How to suppress all autograph warnings from Tensorflow?
- How to spot a greedy algorithm?
- How to spread processes over time getting minimum number of collisions

DSA Fundamentals
Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.
View the 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.