TensorFlow SavedModel is not just a serialization format; it’s the fundamental unit of production for TensorFlow models.
Let’s see a SavedModel in action. Imagine we’ve trained a simple Keras model to classify MNIST digits:
import tensorflow as tf
from tensorflow import keras
# Define a simple model
model = keras.Sequential([
keras.layers.Flatten(input_shape=(28, 28)),
keras.layers.Dense(128, activation='relu'),
keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
# Load and preprocess MNIST data
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train / 255.0
x_test = x_test / 255.0
# Train the model (briefly for demonstration)
model.fit(x_train, y_train, epochs=1)
# Save the model
export_path = '/tmp/my_mnist_model/1'
tf.saved_model.save(model, export_path)
print(f"Model saved to: {export_path}")
This code snippet trains a basic MNIST classifier and then saves it using tf.saved_model.save. The export_path is crucial; it’s a directory, and TensorFlow expects a versioned subdirectory (like /1, /2, etc.) within it. This versioning is key for managing model updates in production.
Now, how do you actually use this saved model? You can load it back into TensorFlow for inference, or more importantly, serve it via TensorFlow Serving.
# Load the saved model
loaded_model = tf.saved_model.load(export_path)
infer = loaded_model.signatures["serving_default"]
# Prepare some test data
sample_image = x_test[0:1] # Take the first test image, keep batch dimension
prediction = infer(tf.constant(sample_image, dtype=tf.float32))
print("Prediction output:", prediction)
The tf.saved_model.load function reconstructs the TensorFlow graph and variables. The signatures attribute is where the magic happens for serving. The "serving_default" signature is automatically created for Keras models and provides a convenient entry point for inference. It maps input tensor names to output tensor names.
Internally, a SavedModel is more than just a Python object. It’s a collection of protocol buffers and files:
saved_model.pb: The core file containing the TensorFlow graph definition and signatures.variables/: A directory holding the trained weights (variables). This includesvariables.data-00000-of-00001andvariables.index.assets/: An optional directory for arbitrary files needed by the model, like vocabulary files or custom ops.
The tf.saved_model.save function handles serializing the Python object’s state (variables) and its computation graph into this structured format. It traces the Python code to build the graph and captures the values of all trainable and non-trainable variables.
The primary problem this solves is decoupling model training from model serving. You train your model in Python (or any TensorFlow-compatible environment), save it as a SavedModel, and then deploy that SavedModel to a serving infrastructure like TensorFlow Serving. This allows you to use different languages or environments for serving (e.g., C++, Java, Go) without needing the original Python training code.
TensorFlow Serving is built to efficiently load and serve SavedModels. It can load multiple versions of a model simultaneously, allowing for A/B testing or gradual rollouts. When a request comes in, TensorFlow Serving executes the graph defined in the SavedModel for the specified signature.
A critical aspect often overlooked is how custom operations or Python functions within your model are handled. If your model relies on Python logic that isn’t directly traceable by TensorFlow’s graph construction, you might need to use tf.function with specific input signatures or ensure that any custom ops are registered and available in the serving environment. For example, if you have a custom Keras layer that uses arbitrary Python code, you’d typically need to use tf.py_function or implement a custom op for deployment.
The next step is typically exploring how to deploy this SavedModel using TensorFlow Serving for high-throughput, low-latency predictions.