Using the SavedModel format in TensorFlow

Tensorflow has a feature of saving model which can be really helpful. It allows you to save your trained neural network for future use and in case you want to use it for deployment. The progress of a model can be saved during the training and even after the training. It is an amazing feature that also lets us use our saved model later without actually writing the whole code again and we can start exactly from where we left off. It helps us to save a lot of resources and time.

Models in TensorFlow can be saved in two different formats i.e SavedModel format and HDF5. The SavedModel format is the default file format in TensorFlow 2.x. These models can also be used for deployment and can be loaded in TensorFlow.js so that we can run them on web browsers as well.

In this attempt, we’ll look at how we can save and use the SavedModel format in Tensorflow. These models can be loaded using tf.keras.models.load_model method.

 

Training and saving a model

The first step is to install and import the TensorFlow module in Python program as you can see below:

import tensorflow as tf
from tensorflow import keras

For training our model we’ll be using the MNIST dataset. MNIST dataset is already packed with training-testing data and respective labels. Next, we’ll be reshaping the images and normalizing the data by changing the range from (0 to 255) to (0 to 1).

(train_images, train_labels), (test_images, test_labels) =   tf.keras.datasets.mnist.load_data()

train_labels = train_labels[:1000]
test_labels = test_labels[:1000]

train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0

Now that we have perfectly split data ready for training the next step is to create a model. For building our model we’ll be using Keras’s Sequential API which allows us to build a model by using sequential layers.

model = tf.keras.models.Sequential([
    keras.layers.Dense(512, activation='relu', input_shape=(784,)),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10)
  ])

model.compile(optimizer='adam',
              loss=tf.losses.SparseCategoricalCrossentropy,
              metrics=['accuracy'])

model.summary()

 

After compiling the model, we have to fit the data and train it on our training data and we will be using test data as our validation data. It’s always a good practice to provide a validation set to check how well the model is performing on unseen data. The validation accuracy can be used to detect if there’s any overfitting.

model.fit(train_images, 
          train_labels,  
          epochs=10,
          validation_data=(test_images,test_labels)

 

Saving the model

model.saveis used to save a model’s architeture, weights and configuration in a file format. The two common model formats are SavedModel and HDF5. Though SavedModel is the default file format in TensorFlow 2 and above versions.

A model can be saved in SavedModel format and can be later used by loading it using tf.keras.models.load_modeland can be later deployed using TensorFlow serving.

After you’ve completed training the model we can save it.

model.save('model_folder/my_model')

A structure of a SavedModel has assets, variables and a saved_model.pb file. The .pb is the MetaGraphDef which has the graph structure saved. The variables folder holds the saved weights. The assets folder allow you to add external files .

If you’re not using Keras then the SavedModel can be very useful. It has an edge over the traditional  tf.train.Saver() function.

tf.train.Saver()only saves the variable but this isn’t the case with SavedModel as it saves the variable in checkpoint files.

You can also check the official documentation for using tf.saved_model.load:

Documentation for tf.saved_model.load

Documentation for custom saved model.