How to load a SavedModel in Keras with TFSMLayer

Keras 3 does not load TensorFlow SavedModel directories with keras.models.load_model(). When the artifact should be reused for inference inside a Keras model, wrap the exported endpoint with keras.layers.TFSMLayer.

TFSMLayer creates a normal Keras layer around a TensorFlow SavedModel or Keras export artifact. It can call the exported inference endpoint, and it can sit inside a larger Keras model when only the SavedModel boundary needs to be preserved.

Use the endpoint name that exists in the artifact. Keras model.export() uses serve by default, while older TensorFlow SavedModel artifacts often expose serving_default.

1. Install Keras and TensorFlow.

   $ python -m pip install --upgrade keras tensorflow

   TFSMLayer is a TensorFlow-specific Keras layer. Use a TensorFlow backend environment for this workflow.

2. Select the TensorFlow backend before importing Keras.

   $ export KERAS_BACKEND=tensorflow

3. Create load_tfsm_layer.py.

   load_tfsm_layer.py

   import os
   import shutil
   from pathlib import Path
    
   os.environ["KERAS_BACKEND"] = "tensorflow"
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import keras
   import numpy as np
    
    
   keras.utils.set_random_seed(21)
    
   export_dir = Path("score_savedmodel")
   if export_dir.exists():
       shutil.rmtree(export_dir)
    
   model = keras.Sequential(
       [
           keras.Input(shape=(3,), name="features"),
           keras.layers.Dense(4, activation="relu", name="hidden"),
           keras.layers.Dense(1, activation="sigmoid", name="score"),
       ]
   )
   model.compile(optimizer="adam", loss="binary_crossentropy")
   model(np.zeros((1, 3), dtype="float32"))
   model.export(export_dir, verbose=False)
    
   try:
       keras.models.load_model(export_dir)
   except ValueError as exc:
       load_model_error = str(exc).split(":", 1)[0]
   else:
       load_model_error = "load_model unexpectedly succeeded"
    
   layer = keras.layers.TFSMLayer(export_dir, call_endpoint="serve")
   inputs = np.array([[0.2, 0.6, 0.4], [0.9, 0.1, 0.7]], dtype="float32")
   outputs = layer(inputs)
   output_array = np.asarray(outputs)
    
   print(f"backend: {keras.backend.backend()}")
   print(f"export directory: {export_dir}")
   print(f"load_model result: {load_model_error}")
   print(f"reloaded layer: {layer.__class__.__name__}")
   print("call endpoint: serve")
   print(f"output shape: {output_array.shape}")
   print(f"first score: {output_array[0, 0]:.6f}")

   Use call_endpoint="serving_default" instead when the SavedModel was not created with model.export() and exposes that endpoint name.

4. Run the script.

   $ python load_tfsm_layer.py
   backend: tensorflow
   export directory: score_savedmodel
   load_model result: File format not supported
   reloaded layer: TFSMLayer
   call endpoint: serve
   output shape: (2, 1)
   first score: 0.646832

   The load_model failure is expected for a SavedModel directory in Keras 3. Confirm the TFSMLayer output shape before wrapping the artifact in a larger model.

5. Use the reloaded layer inside another Keras model when the SavedModel is only one stage of the inference graph.

   inputs = keras.Input(shape=(3,), name="features")
   score = keras.layers.TFSMLayer("score_savedmodel", call_endpoint="serve")(inputs)
   wrapped_model = keras.Model(inputs, score)