A.I, Data and Software Engineering

Save, restore, visualise Graph with TensorFlow v2.0 & KERAS


TensorFlow 2.0 is coming really soon. Therefore, we quickly show some useful features, i.e., save and load a pre-trained model, with v.2 syntax. To make it more intuitive, we will also visualise the graph of the neural network model.

Benefits of saving a model

Quick answer: to save time, easy-share, and fast deploy.

A SavedModel contains a complete TensorFlow program, including weights and computation. It does not require the original model building code to run, which makes it useful for sharing or deploying with different platforms, e.g. TFLite(mobile / IoT), TensorFlow.js (Browsers), TensorFlow Serving (servers)…

With TensorFlow and Keras, we can easily save and restore models, custom models, and sessions. The basic steps are:

  • Create a model
  • Train the model
  • Save the model
  • Share and restore to use.

To demonstrate we will quickly create a sequential neural network using Keras and MNIST fashion dataset. You can try with CIFAR dataset as in this article.

Create a model with Keras and MNIST dataset

Import libraries and enable TensorFlow 2.0

Output: TensorFlow 2.x selected.

Import fashion MNIST

We also create train sets and validation sets.

From the training and validation sets, we can reformat them to tf.data.Dataset by using the two helper methods:

Create and build a Keras sequential model

To visualise the model later, we define the Keras callback, in which ‘logdir’ is where we store the graph info.

Train the model


Before saving the model, let’s see its current description using “model.summary()“:

Save the model

Note that save_format: Either ‘tf’ or ‘h5’, indicating whether to save the model
to Tensorflow SavedModel or HDF5. The default is currently ‘h5’ in TensorFlow 1.*, but it is now ‘tf’ in TensorFlow 2.0.

We save the model in the ‘pretrain’ directory. Now, let check the saved model using ‘saved_model_cli’


What is in the ‘pretrain’ folder:

We have 2 folders, “assets” and “variables”, and one file “save_model.pb”.

Restore the saved model

Loading a model from the ‘pretrain’ directory is as simple as saving the model.

Get the predicted result from the restored model:


Visualise the model with Tensorboard

You can use TensorBoard to visualize your TensorFlow graph, plot quantitative metrics about the execution of your graph, and show additional data like images that pass through it. I find it so powerful and really enjoyable to play with.

If you added the Keras callback as mentioned in the previous section, you will be able to use the Tensorboard embedded to Jupyter notebook.

Visualise the Keras model with Tensorboard
Visualise the model with Tensorboard

And if you want to see the accuracy and loss graph, you can switch to the “SCALARS” tab.

Visualise Train/Loss of Keras model with TensorBoard
Visualise Train/Loss with TensorBoard


In this article, we have demonstrated how easy to save, load, and visualise a model with Keras and TensorBoard. We did not focus on perfecting the model as it was for demo purposes. There are also several changes in TensorFlow v2 that we have not mentioned in this article but may cover some of the most exciting parts in the future posts.

1 comment

A.I, Data and Software Engineering

PetaMinds focuses on developing the coolest topics in data science, A.I, and programming, and make them so digestible for everyone to learn and create amazing applications in a short time.

Pin It on Pinterest


You have successfully subscribed to the newsletter

There was an error while trying to send your request. Please try again.

Petaminds will use the information you provide on this form to be in touch with you and to provide updates.
%d bloggers like this: