TensorFlow
machine learning
neural networks
model deployment
AI inference

TensorFlow 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

TensorFlow inference is the process of using a trained model to produce predictions on new data. The important practical issues are not just loading the model, but making sure preprocessing matches training, the output is interpreted correctly, and the runtime path fits the deployment environment.

What Inference Is

Training updates weights. Inference uses fixed weights.

A normal TensorFlow inference path does four things:

  • load the trained model
  • prepare input data in the expected shape and scale
  • run a forward pass
  • decode the output into something meaningful

If any of those steps mismatches the original training assumptions, the code can run successfully while still producing bad predictions.

A Basic Keras Inference Example

python
1import tensorflow as tf
2import numpy as np
3
4model = tf.keras.Sequential([
5    tf.keras.layers.Input(shape=(4,)),
6    tf.keras.layers.Dense(8, activation="relu"),
7    tf.keras.layers.Dense(3, activation="softmax")
8])
9
10x = np.array([[0.1, 0.2, 0.3, 0.4]], dtype="float32")
11preds = model(x, training=False)
12print(preds.numpy())

The explicit training=False matters for layers such as dropout or batch normalization, where training and inference behavior differ.

Loading a Saved Model for Inference

A common deployment pattern is to save the model after training and reload it later.

python
1import tensorflow as tf
2import numpy as np
3
4model = tf.keras.Sequential([
5    tf.keras.layers.Input(shape=(2,)),
6    tf.keras.layers.Dense(1)
7])
8model.save("/tmp/simple_model")
9
10loaded = tf.keras.models.load_model("/tmp/simple_model")
11result = loaded(np.array([[1.0, 2.0]], dtype="float32"), training=False)
12print(result.numpy())

This is enough for many Python-based applications.

Preprocessing Is Part of Inference

A model does not infer on abstract "data." It infers on tensors shaped and scaled exactly as it expects.

For example, an image model may need:

  • resize to a fixed width and height
  • cast to float32
  • normalize pixel values
  • add a batch dimension
python
1import tensorflow as tf
2
3image = tf.random.uniform((224, 224, 3), maxval=255, dtype=tf.int32)
4image = tf.cast(image, tf.float32) / 255.0
5image = tf.expand_dims(image, axis=0)
6print(image.shape)

If you trained on normalized images and infer on raw 0 to 255 values, the model may look "broken" even though the inference code is technically valid.

Output Interpretation Matters Too

Different tasks produce different output types.

Examples:

  • regression often returns raw numeric values
  • binary classification often returns one sigmoid probability
  • multiclass classification often returns a vector of class scores or probabilities
  • detection and segmentation models return structured multi-tensor outputs

That means inference code should include postprocessing, not just predict.

python
1import numpy as np
2
3probs = np.array([[0.1, 0.7, 0.2]])
4label = int(np.argmax(probs, axis=1)[0])
5print(label)

Without correct output decoding, a valid model run is not yet a valid application prediction.

Batch Inference Versus Single-Item Inference

TensorFlow models usually accept batches. Even if you only need one prediction, the data is often shaped as a batch of size 1.

Batching matters operationally too:

  • single-item inference is simple and low-latency
  • batched inference improves throughput
  • very large batches may increase memory pressure or latency

The right choice depends on whether your deployment cares more about per-request latency or total throughput.

Common Pitfalls

The most common mistake is forgetting to match training-time preprocessing during inference.

Another mistake is not setting inference mode explicitly when model behavior differs between training and inference.

A third issue is interpreting raw outputs incorrectly, such as taking logits as probabilities or reading class indices from the wrong axis.

Finally, deployment bugs often come from shape mismatches. A model trained on (batch, 224, 224, 3) will not accept an unbatched (224, 224, 3) array in every pipeline without adjustment.

Summary

  • TensorFlow inference means loading a trained model and running a forward pass on new data.
  • Preprocessing must match what the model saw during training.
  • Use training=False for layers with different inference behavior.
  • Interpret outputs according to the task, not just the tensor shape.
  • Batch size affects both latency and throughput.
  • Most inference bugs are data-contract bugs, not TensorFlow runtime bugs.

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

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.