TensorFlow SavedModel: Converting Keras Models to SavedModel

TensorFlow's SavedModel format is a universal serialization format meant to save TensorFlow models for use in other applications. As a developer, you might want to convert a Keras model to a TensorFlow SavedModel for easy deployment, inference, or further training in different environments. This article will guide you through the process of converting a Keras model to a SavedModel using TensorFlow.

Understanding TensorFlow SavedModel

The SavedModel format contains both metadata and the intermediates (such as graph states and weights) needed to restore the TensorFlow model. It's robust, allowing you to deploy models in various environments with consistency and accuracy.

Creating and Training a Keras Model

Let's start by creating and training a simple Keras model, which we will then convert to a SavedModel.

from tensorflow import keras
from tensorflow.keras import layers

# Simple Sequential Model
model = keras.Sequential([
    layers.Dense(64, activation='relu', input_shape=(32,)),
    layers.Dense(64, activation='relu'),
    layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Dummy Data
import numpy as np

data = np.random.random((1000, 32))
labels = np.random.randint(10, size=(1000, 1))

# Training the Model
model.fit(data, labels, epochs=10, batch_size=32)

Saving the Keras Model as a SavedModel

Once your Keras model is trained, the next step is to convert and export it as a SavedModel. TensorFlow provides an easy way to achieve this.

import os

# Define the path for saving the model
saved_model_path = "./my_model"

# Save the model
model.save(saved_model_path, save_format='tf')

This command will save the model in a specified directory (in this case ./my_model), generating files and directories that conform to the SavedModel format. The resulting structure includes:

• saved_model.pb or saved_model.pbtxt: Contains the model graph and metadata.
• variables/: Contains subdirectories for each of the model's Variable objects.

Loading the SavedModel

To ensure everything has been saved correctly and that you can reload the SavedModel format, TensorFlow offers an easy way to reload the model.

import tensorflow as tf

# Load the saved model
loaded_model = tf.keras.models.load_model(saved_model_path)

# Use the loaded model for predictions
dummy_input = np.random.random((1, 32))
output = loaded_model.predict(dummy_input)
print("Prediction:", output)

This code loads the SavedModel and allows you to use it for making predictions, demonstrating that the model and weights have been successfully restored.

Advantages of Using SavedModel

The SavedModel format is crucial for several reasons:

• Portability: SavedModel is language-agnostic and works well with different versions of TensorFlow.
• Flexibility: Compatible with TensorFlow Serving, TensorFlow Lite, and TensorFlow.js, enabling deployment across different platforms (e.g., web, mobile).
• Robustness: Supports complex variable deserialisation and optimisation, helping in fine-tuning models further without barriers.

Conclusion

Converting a Keras model to TensorFlow's SavedModel is a straightforward process that enhances the interoperability and deployment capabilities across various platforms. Now that you're more familiar with the steps, you should have a solid foundation that allows you to leverage this format effectively in your machine learning projects.