How to save and load models in TensorFlow

Note: This article assumes you have a working knowledge of machine learning and the Tensorflow and Keras framework.

If you just want to get right into saving and loading models, you can skip the next 3 paragraphs

Statistics have shown that more than 80% of machine learning models do not make it to production, that is, they do not solve any business problem(s) but instead remain in a Jupyter notebook or random GitHub repo, etc.

For a machine learning model to provide value, it has to be made accessible to the world or the specific set of people who will use it. These people are mostly non-technical; even if they are, they want to accomplish their set objectives in the easiest way possible. Hence the need to make these models as easily accessible as possible for them to be useful.

In other, for models to be deployed, they first have to be saved during or after training and thereafter loaded into the environment it will be used to perform predictions.

Saving & loading TensorFlow models

TensorFlow models can be saved in two different file formats; SavedModel and HDF5.

SavedModel Format

This is the default file format for TensorFlow. It contains a complete TensorFlow program, including trained parameters and computation. It does not require the original model building code to run, which makes it useful for sharing or deployment. You can save and load a model in the SavedModel format using the following APIs:

Low-level API

  • Save: tf.saved_model.save(model, path_to_dir)

  • Load: tf.saved_model.load(path_to_dir)

High-level API - Using Keras

  • Save: tf.keras.models.save_model( model,filepath, verwrite=True, include_optimizer=True, save_format=None, signatures=None, options=None, save_traces=True)

  • Load: tf.keras.models.load_model('path_to_dir')

HDF5 Format

Hierarchical Data Format (HDF) is a set of file formats (HDF4, HDF5) designed to store and organize large amounts of data.

  • Save: model.save("my_h5_model.h5")

  • Load: keras.models.load_model('path_to_dir')

The difference between HDF5 and SavedModel is that HDF5 uses object configs to save the model architecture, while SavedModel saves the execution graph. Thus, SavedModels are able to save custom objects like subclassed models and custom layers without requiring the original code.

Source: TensorFlow Website