TensorFlow How to predict from a SavedModel?
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
SavedModel is TensorFlow's standard export format for serving and reuse. If you can load the model directory, inspect its signatures, and feed tensors with the expected shape and dtype, prediction becomes straightforward.
Load the Model and Inspect Its Signatures
The most important detail is that a SavedModel is not just weights on disk. It also contains callable functions, often exposed under a serving_default signature. When you load the directory with tf.saved_model.load, TensorFlow returns an object with those exported functions attached.
structured_input_signature tells you the input names, shapes, and dtypes. That matters because prediction errors usually come from passing NumPy arrays or tensors that do not match the exported signature.
If you are exporting a model yourself, this small example creates a simple network and saves it in a way that is easy to reload later:
After export, the directory contains metadata, variables, and any assets the model needs at inference time.
Run Prediction with the Exported Signature
Once the model is loaded, call the signature like a regular TensorFlow function. The key is to pass keyword arguments that match the exported input name. In many Keras exports, the input name is something generic such as keras_tensor, so inspecting the signature first is worth the effort.
The return value is usually a dictionary mapping output names to tensors. You can convert the tensor to NumPy if you are working in eager mode:
If you do not know the argument name ahead of time, inspect it dynamically:
This approach is especially useful when the model was exported by another team or generated in a notebook months ago.
Use the Keras Loading Path When Appropriate
Some TensorFlow models are better loaded with Keras APIs instead of the low-level SavedModel interface. If the export was created with model.save() in a format that Keras understands, tf.keras.models.load_model gives you back an object with the familiar predict method.
The rule of thumb is simple:
- Use
tf.saved_model.loadwhen you are consuming a generic TensorFlow export. - Use
tf.keras.models.load_modelwhen you are loading a Keras model and want the higher-level API back.
Common Pitfalls
The most common failure is sending the wrong input name. SavedModel functions are strict about keyword names, so guessing often produces an error even when the tensor shape is correct.
Another common issue is dtype mismatch. A model trained on float32 inputs may reject int64 arrays or silently force conversions that hurt performance. Create tensors with an explicit dtype during inference.
Shape mismatches also show up frequently. If the signature expects a batch shaped like (None, 2), a one-dimensional tensor such as [5.0, 1.0] is not enough. Wrap the sample in an outer batch dimension.
Developers also confuse output names. A result dictionary may use names such as output_0 or a layer-derived label rather than predictions. Print structured_outputs once and use the actual key.
Finally, do not assume every export has a serving_default signature. Some models define custom signatures, so list the available keys before calling anything.
Summary
- '
tf.saved_model.loadis the standard way to load a TensorFlowSavedModeldirectory.' - Inspect signatures first so you know the exact input names, shapes, and dtypes.
- Call the prediction function with keyword arguments that match the exported signature.
- Expect a dictionary of output tensors and inspect the real output keys.
- Use the Keras loading path only when the model was saved in a Keras-compatible format.
Related reading
- TensorFlow How to predict from a SavedModel?
- Tensorflow How to randomly crop input images and labels in the same way?
- Tensorflow How to reduce memory footprint for inference only models?
- Tensorflow How to replace a node in a calculation graph?
- Tensorflow How to replace or modify gradient?
- tensorflow how to rotate an image for data augmentation?
- TensorFlow how to safely terminate training manually KeyboardInterrupt
- Tensorflow How to switch channels of a tensor from RGB to BGR?
.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.