Tensorflow
Keras
Model Optimization
Inference Speed
Deep Learning

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.

Practice ML system design

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.

python
1import time
2import tensorflow as tf
3
4model = tf.keras.applications.MobileNetV2(weights=None, classes=10)
5sample = tf.random.normal((1, 224, 224, 3))
6
7# Warm up graph creation and kernel selection.
8for _ in range(10):
9    _ = model(sample, training=False)
10
11start = time.perf_counter()
12for _ in range(100):
13    _ = model(sample, training=False)
14elapsed = time.perf_counter() - start
15
16print(f"Average latency: {elapsed / 100:.6f} seconds")

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.

python
1import tensorflow as tf
2
3@tf.function
4def infer(batch):
5    return model(batch, training=False)
6
7sample = tf.random.normal((8, 224, 224, 3))
8predictions = infer(sample)
9print(predictions.shape)

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.

python
1batches = [
2    tf.random.normal((1, 224, 224, 3)),
3    tf.random.normal((8, 224, 224, 3)),
4    tf.random.normal((32, 224, 224, 3)),
5]
6
7for batch in batches:
8    start = time.perf_counter()
9    _ = infer(batch)
10    elapsed = time.perf_counter() - start
11    print(batch.shape[0], elapsed)

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:

python
1converter = tf.lite.TFLiteConverter.from_keras_model(model)
2converter.optimizations = [tf.lite.Optimize.DEFAULT]
3tflite_model = converter.convert()
4
5with open("model.tflite", "wb") as f:
6    f.write(tflite_model)

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.

python
1import tensorflow as tf
2
3tf.config.optimizer.set_jit(True)
4
5@tf.function
6def fast_infer(batch):
7    return model(batch, training=False)

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:

python
from tensorflow.keras import mixed_precision

mixed_precision.set_global_policy("mixed_float16")

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 direct model(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
Course
Intermediate
27 lessons
15 hours
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 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.