Deploy Keras Models with TensorFlow SavedModel

Problem and takeaway

Production inference breaks when a Keras model is saved as a single HDF5 file: TensorFlow Serving expects a directory, custom layer details can be lost, and overwriting a live model file exposes readers to partial writes. The practical fix is to save inference artifacts as TensorFlow SavedModel directories, write them to temporary storage, and move them into versioned deployment paths.

Desired outcome

A trained Keras model is serialized to a SavedModel directory, reloaded with tf.saved_model.load, and used to run inference on sample input. A versioned folder layout lets serving traffic switch to the new model only after the directory is completely written.

Prerequisites

• TensorFlow 2 installed in the Python environment used for training and serving.
• A Keras model object in memory, or a script that produces one.
• Write permission on the directory that will hold model artifacts.
• Sample input data matching the model's expected shape and dtype.

Confirm the installed major version before applying version-specific save behavior:

python -c "import tensorflow as tf; print(tf.__version__)"

Run this in the shell where TensorFlow is installed. No elevated permissions are required. Expected output begins with 2.. If it fails, activate the correct environment or install TensorFlow before continuing.

Save the model as a SavedModel directory

SavedModel is TensorFlow 2's directory-based format. It bundles the computation graph, variables, and metadata in files such as saved_model.pb and a variables/ subdirectory. HDF5 .h5 files remain supported but are less suitable for production serving and can omit TensorFlow-specific layer details.

For inference-only artifacts, disable optimizer state to keep the directory smaller and avoid serializing training-only information:

import tensorflow as tf

# model: a tf.keras.Model instance
export_dir = "/tmp/keras_saved_model"

tf.keras.models.save_model(
    model,
    export_dir,
    save_format="tf",
    include_optimizer=False,
)

Run this Python code in the environment that holds the trained model. export_dir is a placeholder; replace it with a writable path. If the directory already exists, remove it or choose a new version folder first, because overwriting an existing SavedModel directory can leave stale files depending on the TensorFlow version.

Format comparison

Aspect	SavedModel directory	HDF5 .h5 file
Default format in TensorFlow 2	Yes	No
Structure	Directory with saved_model.pb and variables/	Single file
Suitable for TensorFlow Serving	Yes	No
Optimizer state	Controllable with include_optimizer	Included
Partial-write risk during deployment	Reduced with atomic directory move	Higher when overwriting a live file

Load and run inference

Reload the directory with the SavedModel API rather than the Keras HDF5 loader. This reduces dependence on the original Python class file:

import tensorflow as tf
import numpy as np

loaded = tf.saved_model.load("/tmp/keras_saved_model")
infer = loaded.signatures["serving_default"]

sample = np.random.rand(1, 28, 28, 1).astype("float32")
prediction = infer(tf.constant(sample))
print(prediction)

The signatures dictionary exposes concrete inference functions. Confirm the signature key name before relying on it in serving code; a missing key raises an error at load time.

Versioned deployment with an atomic move

Serving systems can read a model directory while it is in use. Writing directly into a live deployment path can expose a partially saved model. Write to a temporary location first, then move the complete directory into a versioned folder:

tmp_dir="/tmp/model_build/mymodel"
final_dir="/deploy/models/mymodel/2"

mv "$tmp_dir" "$final_dir"

Run mv in a shell with write permission on /deploy/models/. On the same filesystem, a move is effectively atomic: readers see either the previous version or the complete new directory, never a half-written one. On a network filesystem, verify that the move completed and list the directory contents before removing the previous version.

Custom layers, losses, and metrics

If the model uses a custom Keras layer, loss, or metric, the SavedModel directory stores the graph, but the Keras load API may still need the Python class to reconstruct the object. Pass it through custom_objects:

tf.keras.models.load_model(
    "/tmp/keras_saved_model",
    custom_objects={"MyLayer": MyLayer},
)

Prefer tf.saved_model.load for inference-only workloads because it depends less on the original class definitions. If a custom object is missing at load time, the loader raises an error before inference can run. Import or register the class, then retry.

Expected checks

1. List the export directory and confirm saved_model.pb and variables/ exist.
2. Reload with tf.saved_model.load and run a sample batch through the restored concrete function.
3. Compare original and reloaded predictions on the same input with numpy.allclose; outputs should match within tolerance.
4. Load a model containing a custom layer without custom_objects to confirm the expected registration error, then add the class and verify recovery.

Limitations

• SavedModel directories consume more space than a single .h5 file because weights are stored as separate checkpoint files.
• Large writes to network filesystems can fail partway; verify directory contents after save before deleting the previous model.
• Format mismatch causes load failures: do not load an HDF5 file with SavedModel APIs or a SavedModel directory with HDF5 loaders.
• Behavior can vary between TensorFlow 2 minor versions; confirm the installed version before standardizing on signature keys or save options.

Rollback and recovery

Saving changes state by writing files. Keep the previous version directory in place until the new one passes the checks above. If the new model fails verification, remove the failed directory and point the serving configuration back to the previous version folder. Do not overwrite a live model directory in place.