How to load a TensorFlow SavedModel

A SavedModel is the TensorFlow handoff artifact for running an exported inference graph outside the training script. Loading it in Python is a fast smoke test before the model moves into a service, batch job, converter, or another runtime.

The low-level tf.saved_model.load() API returns a trackable object, not a rehydrated training model. Keras 3 exports commonly expose serve and serving_default endpoints, and the selected endpoint controls the input key and output dictionary used by the caller.

A complete export directory must be available in the active environment before loading begins. Copy the whole directory, including saved_model.pb and variables, and prepare a sample tensor with the shape and dtype shown by the serving signature.

1. Open a terminal in the Python environment that imports TensorFlow and can read the SavedModel directory.

   Use tf.saved_model.load() for direct TensorFlow endpoint calls. Use a Keras TFSMLayer when the SavedModel has to become a layer inside a Keras model.

2. Confirm the export path points at the complete SavedModel directory.

   $ ls support_score_savedmodel
   assets
   fingerprint.pb
   saved_model.pb
   variables

   Copy the whole directory, not only saved_model.pb, because variables and assets can live outside the protobuf file.

3. Save the loader script as

   load_savedmodel_demo.py

   .

   load_savedmodel_demo.py

   import argparse
   import os
   import pathlib
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
    
    
   def shape_tuple(shape):
       dims = shape.as_list()
       return tuple(None if dim is None else dim for dim in dims)
    
    
   parser = argparse.ArgumentParser()
   parser.add_argument("savedmodel_dir", help="Path to a TensorFlow SavedModel directory")
   args = parser.parse_args()
    
   model_dir = pathlib.Path(args.savedmodel_dir)
   if not tf.saved_model.contains_saved_model(str(model_dir)):
       raise SystemExit(f"{model_dir} is not a TensorFlow SavedModel directory")
    
   artifact = tf.saved_model.load(str(model_dir))
   signature_names = sorted(artifact.signatures.keys())
   if not signature_names:
       raise SystemExit("No callable signatures found in SavedModel")
    
   signature_name = "serving_default" if "serving_default" in artifact.signatures else signature_names[0]
   signature = artifact.signatures[signature_name]
   _, keyword_inputs = signature.structured_input_signature
   if len(keyword_inputs) != 1:
       raise SystemExit(f"Expected one keyword input, found {sorted(keyword_inputs.keys())}")
    
   input_name, input_spec = next(iter(keyword_inputs.items()))
   sample_batch = tf.constant(
       [
           [0.15, 0.75, 0.35, 0.45],
           [0.85, 0.25, 0.65, 0.75],
       ],
       dtype=input_spec.dtype,
   )
   outputs = signature(**{input_name: sample_batch})
    
   print(f"TensorFlow {tf.__version__}")
   print(f"SavedModel dir: {model_dir}")
   print(f"Contains SavedModel: {tf.saved_model.contains_saved_model(str(model_dir))}")
   print(f"Available signatures: {signature_names}")
   print(f"Selected signature: {signature_name}")
   print(f"Input tensor: {input_name} shape={shape_tuple(input_spec.shape)} dtype={input_spec.dtype.name}")
   print(f"Output keys: {sorted(outputs.keys())}")
    
   for output_name in sorted(outputs.keys()):
       output_tensor = outputs[output_name]
       rounded = tf.round(output_tensor * 10000) / 10000
       print(f"Output {output_name} shape: {tuple(output_tensor.shape)}")
       print(f"Output {output_name} values:")
       print(rounded.numpy())

   Replace support_score_savedmodel and the sample_batch values with the directory, shape, and dtype from the exported model.

4. Run the loader against the SavedModel directory.

   $ python3 load_savedmodel_demo.py support_score_savedmodel
   TensorFlow 2.21.0
   SavedModel dir: support_score_savedmodel
   Contains SavedModel: True
   Available signatures: ['serve', 'serving_default']
   Selected signature: serving_default
   Input tensor: features shape=(None, 4) dtype=float32
   Output keys: ['output_0']
   Output output_0 shape: (2, 1)
   Output output_0 values:
   [[0.54  ]
    [0.5525]]

   The printed Input tensor and Output keys are the values to map into a serving request, batch inference job, or downstream client.

5. Remove the temporary smoke-test script after recording the signature and output contract.

   $ rm load_savedmodel_demo.py

   Skip this cleanup when the loader becomes part of a release checklist or model handoff test.