The world’s leading publication for data science, AI, and ML professionals.

How to use a saved model in Tensorflow 2.x

A tutorial on saving and reusing Tensorflow-trained models

Image generated by the author using wordart.com
Image generated by the author using wordart.com

In my previous article, I wrote about model validation, regularization, and callbacks using TensorFlow 2.x. In the machine-learning pipeline, creating a trained model is not enough. What are we going to do with the trained model once we have finished the training, validated, and tested it with a portion of data set aside? In practice, we would like to import such a trained model so that it can be of use in some practical applications. For example, let’s say I trained a model on camera images to recognize pedestrians. Ultimately, I want to use the trained model to make a real-time prediction of detecting pedestrians with a camera mounted on a self-driving car. Additionally, training a model also needs saving the model as checkpoints, especially when you are training a model on a really large dataset or training time is in the order of hours. Model saving is also useful in case your training gets interrupted for some reasons such as a flaw in your programming logic, the battery of your laptop died, there was an I/O error, etc.

Tensorflow 2: Model validation, regularization, and callbacks

Ultimately, I want to use the trained model to make a real-time prediction of detecting pedestrians with a camera mounted on a self-driving car.

Saving and loading models

There are a couple of things we can do while saving a model. Do we want to save the model weights and training parameters in every iteration (epochs), every once in a while, or once training has finished? We can use built-in callbacks just like we saw in my previous article to automatically save the model weights during the training process. Alternatively, we can also save the model weights and other necessary information once training has finished.

There are two main formats for saved models: One in native TensorFlow, and the other in HDF5 format since we are using TensorFlow through Keras API.

An example of saving the model during the training procedure:

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.losses import BinaryCrossentropy
from tensorflow.keras.callbacks import ModelCheckpoint
model = Sequential( [
Dense(128, activation='sigmoid', input_shape = (10, )),
Dense(1)])
model.compile(optimizer = 'sgd', loss = BinaryCrossentropy(from_logits  = True))
checkpoint = ModelCheckpoint('saved_modelname', save_weights_only=True)
model.fit(X_train, y_train, epochs = 10, callbacks = [checkpoint])

You can see that, I used ModelCheckpoint class to create an object checkpoint that takes an argument which will be used as a filename to save the model. Since save_weights_only=True is used, only weights will be saved, and network architecture will not be saved. Finally, we pass callback = [checkpoint] to the fit function.

If instead of 'saved_modelname' if we supply 'saved_modelname.h5 , then the model will be saved in HDF5 format.

To load the weights of previously saved, we call load_weights function.

model = Sequential( [
Dense(128, activation='sigmoid', input_shape = (10, )),
Dense(1)])
model.load_weights('saved_modelname')

Saving the model weights manually without the callbacks

We can also save the models manually by calling save_weights at the end of the training

model = Sequential( [
Dense(128, activation='sigmoid', input_shape = (10, )),
Dense(1)])
model.compile(optimizer = 'sgd', loss = BinaryCrossentropy(from_logits  = True))
model.fit(X_train, y_train, epochs = 10)
model.save_weights("saved_modelname")

You can also analyze the directory where the model is saved:

total 184K
-rw-r--r-- 1 ivory ivory   61 Jan 12 01:08 saved_modelname
-rw-r--r-- 1 ivory ivory 174K Jan 12 01:08 saved_modelname.data-00000-of-00001
-rw-r--r-- 1 ivory ivory 2.0K Jan 12 01:08 saved_modelname.index

Here, you can see that the actual model saved is saved_modelname.data-00000-of-00001 and the rest of the file contains metadata.

Saving the entire model

Until now, we saw that we were only saving the model weights. However, saving the entire model is very easy. Just pass save_weights_only=False while instantiating ModelCheckpoint class.

checkpoint_dir = 'saved_model'
checkpoint = ModelCheckpoint(filepath=checkpoint_dir, 
                            frequency = "epoch",
                            save_weights_only = False,
                             verbose= True)
model.fit(X_train, y_train, callbacks=[checkpoint])

In this case, a new directory is created with the following content:

total 128
drwxr-xr-x 2 ivory ivory   4096 Jan 12 01:14 assets
-rw-r--r-- 1 ivory ivory 122124 Jan 12 01:14 saved_model.pb
drwxr-xr-x 2 ivory ivory   4096 Jan 12 01:14 variables

In this case, the main model is saved in the file saved_model.pb and other files are metadata.

Finally, we can use the saved model as follows:

from tensorflow.keras.models import load_model
model = load_model(checkpoint_dir)

If we want to save the model once the training procedure is finished, we can call save function as follows:

model.save("mysavedmodel")

If you use model.save("mysavedmodel.h5"), then the model will be saved as a single file mysavedmodel.h5 .

The saved model can be used to make predictions using a brand new data set.

model.predict(X_test)

A more descriptive example is given in my GitHub repo at https://github.com/rahulbhadani/medium.com/blob/master/01_12_2021/Saving_Model_TF2.ipynb.

Acknowledgment

The article is motivated by the author’s learning from the Tensorflow2 Coursera course https://www.coursera.org/learn/getting-started-with-tensor-flow2/. Readers may find similarities in the presented example with examples from the Coursera course. Permission to use the code example (in its original form or with some modification) from the Coursera course was obtained from the instructor.


Related Articles