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.
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.
The most common signature is serving_default, but a SavedModel may contain others.
Choose one and inspect it:
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.
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:
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
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:
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.loadand inspectloaded.signaturesfirst. - Use
structured_input_signatureandstructured_outputsto 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
- Inference using saved model in Tensorflow 2 how to control in/output?
- Initialize keras placeholder as Input to a Custom Layer
- Initializing LSTM hidden state Tensorflow/Keras
- Initializing tensorflow Variable with an array larger than 2GB
- Inference with TensorRT .engine file on python
- Inferring templates from a collection of strings
- Input 0 is incompatible with layer flatten_2 expected min_ndim3, found ndim2
- Input 0 of layer conv1d is incompatible with the layer expected min_ndim3, found ndim2. Full shape received None, 30
.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.