TensorFlow
machine learning
inference
saved model
input output control

Inference using saved model in Tensorflow 2 how to control in/output?

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

In TensorFlow 2, controlling SavedModel inputs and outputs during inference is mostly about signatures. A SavedModel can expose one or more concrete callable functions, each with named inputs and named outputs. If you know how to inspect those signatures, choose the right one, and call it by keyword, you can control inference much more precisely than by treating the model as a black box.

Load The SavedModel And Inspect Signatures

The first step is to load the model and inspect what callable signatures it exports.

python
1import tensorflow as tf
2
3loaded = tf.saved_model.load("saved_model_dir")
4print(list(loaded.signatures.keys()))

The most common signature is serving_default, but a SavedModel may contain others.

Choose one and inspect it:

python
infer = loaded.signatures["serving_default"]
print(infer.structured_input_signature)
print(infer.structured_outputs)

This tells you the input names, input shapes, dtypes, and output names that TensorFlow expects.

Call The Signature By Named Input

If the signature says the input is named image, call it with that exact keyword.

python
1import tensorflow as tf
2
3loaded = tf.saved_model.load("saved_model_dir")
4infer = loaded.signatures["serving_default"]
5
6x = tf.random.uniform((1, 224, 224, 3), dtype=tf.float32)
7outputs = infer(image=x)
8
9print(outputs)
10print(outputs.keys())

The returned value is usually a dictionary-like mapping from output names to tensors.

That output naming is the main way you control which result you read.

Why Names Matter

A common mistake is assuming the SavedModel accepts positional input tensors like a regular Python function you wrote yourself. Signature functions are stricter. They are exported with specific names and shapes.

If the signature expects:

  • 'tokens'
  • 'mask'

then inference should look more like:

python
result = infer(tokens=tokens_tensor, mask=mask_tensor)

not an arbitrary positional call.

The same principle applies to outputs. Read the specific key you want from the returned mapping.

Example With Multiple Inputs And Outputs

python
1import tensorflow as tf
2
3loaded = tf.saved_model.load("saved_model_dir")
4infer = loaded.signatures["serving_default"]
5
6print(infer.structured_input_signature)
7print(infer.structured_outputs)
8
9features = tf.random.uniform((2, 16), dtype=tf.float32)
10metadata = tf.random.uniform((2, 4), dtype=tf.float32)
11
12outputs = infer(features=features, metadata=metadata)
13
14prediction = outputs["prediction"]
15embedding = outputs["embedding"]
16
17print(prediction.shape)
18print(embedding.shape)

This is the core control pattern: inspect names first, then call by those names, then read the named outputs you care about.

Control Inputs And Outputs When Saving

The best time to control inference I/O is often when saving the model. If you export custom signatures with explicit input and output names, the loaded SavedModel becomes much easier to use downstream.

A simplified example:

python
1import tensorflow as tf
2
3class MyModule(tf.Module):
4    @tf.function(input_signature=[tf.TensorSpec([None, 4], tf.float32, name="features")])
5    def serve(self, features):
6        prediction = features * 2.0
7        return {"prediction": prediction}
8
9module = MyModule()
10tf.saved_model.save(module, "saved_model_dir", signatures={"serving_default": module.serve})

Now the exported model has a predictable input name and a predictable output key.

Keras Models And SavedModel

If the SavedModel came from a Keras model, you may also be able to load it with Keras APIs in some workflows. But when you specifically need signature-level control of inference inputs and outputs, tf.saved_model.load plus signature inspection is often the clearest route.

This is especially true in deployment or interoperability scenarios.

Common Pitfalls

  • Guessing input names instead of inspecting structured_input_signature.
  • Assuming the SavedModel returns one unnamed tensor when it actually returns a mapping.
  • Feeding tensors with the wrong dtype or shape for the exported signature.
  • Treating signature functions like arbitrary Python callables with positional arguments.
  • Exporting vague or default signatures and then wondering why downstream inference is hard to control.

Summary

  • In TensorFlow 2 SavedModel inference, input and output control is primarily handled through signatures.
  • Load the model with tf.saved_model.load and inspect loaded.signatures first.
  • Use structured_input_signature and structured_outputs to discover the exact names and shapes.
  • Call the signature with named arguments and read named outputs from the returned mapping.
  • If you control the export step, define clear custom signatures to make inference easier later.

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.