Serialization and saving
Authors: Kathy Wu, Francois Chollet
Date created: 2020/04/28
Last modified: 2020/04/28
Description: Complete guide to saving & serializing models.
A Keras model consists of multiple components:
- The architecture, or configuration, which specifies what layers the model contain, and how they’re connected.
- A set of weights values (the «state of the model»).
- An optimizer (defined by compiling the model).
- A set of losses and metrics (defined by compiling the model or calling add_loss() ).
The Keras API makes it possible to save all of these pieces to disk at once, or to only selectively save some of them:
- Saving everything into a single archive in the TensorFlow SavedModel format (or in the older Keras H5 format). This is the standard practice.
- Saving the architecture / configuration only, typically as a JSON file.
- Saving the weights values only. This is generally used when training the model.
Let’s take a look at each of these options. When would you use one or the other, and how do they work?
How to save and load a model
If you only have 10 seconds to read this guide, here’s what you need to know.
Saving a Keras model:
Loading the model back:
Now, let’s look at the details.
Whole-model saving & loading
You can save an entire model to a single artifact. It will include:
- The model’s architecture/config
- The model’s weight values (which were learned during training)
- The model’s compilation information (if compile() was called)
- The optimizer and its state, if any (this enables you to restart training where you left)
- model.save() or tf.keras.models.save_model()
There are two formats you can use to save an entire model to disk: the TensorFlow SavedModel format, and the older Keras H5 format. The recommended format is SavedModel. It is the default when you use model.save() .
You can switch to the H5 format by:
- Passing save_format=’h5′ to save() .
- Passing a filename that ends in .h5 or .keras to save() .
SavedModel is the more comprehensive save format that saves the model architecture, weights, and the traced Tensorflow subgraphs of the call functions. This enables Keras to restore both built-in layers as well as custom objects.
What the SavedModel contains
Calling model.save(‘my_model’) creates a folder named my_model , containing the following:
The model architecture, and training configuration (including the optimizer, losses, and metrics) are stored in saved_model.pb . The weights are saved in the variables/ directory.
For detailed information on the SavedModel format, see the SavedModel guide (The SavedModel format on disk).
How SavedModel handles custom objects
When saving the model and its layers, the SavedModel format stores the class name, call function, losses, and weights (and the config, if implemented). The call function defines the computation graph of the model/layer.
In the absence of the model/layer config, the call function is used to create a model that exists like the original model which can be trained, evaluated, and used for inference.
Nevertheless, it is always a good practice to define the get_config and from_config methods when writing a custom model or layer class. This allows you to easily update the computation later if needed. See the section about Custom objects for more information.
The first loaded model is loaded using the config and CustomModel class. The second model is loaded by dynamically creating the model class that acts like the original model.
Configuring the SavedModel
New in TensorFlow 2.4 The argument save_traces has been added to model.save , which allows you to toggle SavedModel function tracing. Functions are saved to allow the Keras to re-load custom objects without the original class definitions, so when save_traces=False , all custom objects must have defined get_config / from_config methods. When loading, the custom objects must be passed to the custom_objects argument. save_traces=False reduces the disk space used by the SavedModel and saving time.
Keras H5 format
Keras also supports saving a single HDF5 file containing the model’s architecture, weights values, and compile() information. It is a light-weight alternative to SavedModel.
Keras SavedModel format limitations:
The tracing done by SavedModel to produce the graphs of the layer call functions allows SavedModel be more portable than H5, but it comes with drawbacks.
- Can be slower and bulkier than H5.
- Cannot serialize the ops generated from the mask argument (i.e. if a layer is called with layer(. mask=mask_value) , the mask argument is not saved to SavedModel).
- Does not save the overridden train_step() in subclassed models.
Custom objects that use masks or have a custom training loop can still be saved and loaded from SavedModel, except they must override get_config() / from_config() , and the classes must be passed to the custom_objects argument when loading.
- External losses added via model.add_loss() are not saved (unlike SavedModel). If you have such losses & metrics on your model and you want to resume training, you need to add these losses back yourself after loading the model. Note that this does not apply to losses/metrics created inside layers, e.g. self.add_loss() . As long as the layer gets loaded, these losses & metrics are kept, since they are part of the call method of the layer.
- The computation graph of custom objects such as custom layers is not included in the saved file. At loading time, Keras will need access to the Python classes/functions of these objects in order to reconstruct the model. See Custom objects.
- Does not support preprocessing layers.
Saving the architecture
The model’s configuration (or architecture) specifies what layers the model contains, and how these layers are connected*. If you have the configuration of a model, then the model can be created with a freshly initialized state for the weights and no compilation information.
*Note this only applies to models defined using the functional or Sequential apis not subclassed models.
Configuration of a Sequential model or Functional API model
These types of models are explicit graphs of layers: their configuration is always available in a structured form.
- get_config() and from_config()
- tf.keras.models.model_to_json() and tf.keras.models.model_from_json()
get_config() and from_config()
Calling config = model.get_config() will return a Python dict containing the configuration of the model. The same model can then be reconstructed via Sequential.from_config(config) (for a Sequential model) or Model.from_config(config) (for a Functional API model).
The same workflow also works for any serializable layer.
Sequential model example:
Functional model example:
to_json() and tf.keras.models.model_from_json()
This is similar to get_config / from_config , except it turns the model into a JSON string, which can then be loaded without the original model class. It is also specific to models, it isn’t meant for layers.
Models and layers
The architecture of subclassed models and layers are defined in the methods __init__ and call . They are considered Python bytecode, which cannot be serialized into a JSON-compatible config — you could try serializing the bytecode (e.g. via pickle ), but it’s completely unsafe and means your model cannot be loaded on a different system.
In order to save/load a model with custom-defined layers, or a subclassed model, you should overwrite the get_config and optionally from_config methods. Additionally, you should register the custom object so that Keras is aware of it.
Custom-defined functions (e.g. activation loss or initialization) do not need a get_config method. The function name is sufficient for loading as long as it is registered as a custom object.
Loading the TensorFlow graph only
It’s possible to load the TensorFlow graph generated by the Keras. If you do so, you won’t need to provide any custom_objects . You can do so like this:
Note that this method has several drawbacks: * For traceability reasons, you should always have access to the custom objects that were used. You wouldn’t want to put in production a model that you cannot re-create. * The object returned by tf.saved_model.load isn’t a Keras model. So it’s not as easy to use. For example, you won’t have access to .predict() or .fit()
Even if its use is discouraged, it can help you if you’re in a tight spot, for example, if you lost the code of your custom objects or have issues loading the model with tf.keras.models.load_model() .
Defining the config methods
- get_config should return a JSON-serializable dictionary in order to be compatible with the Keras architecture- and model-saving APIs.
- from_config(config) ( classmethod ) should return a new layer or model object that is created from the config. The default implementation returns cls(**config) .
Registering the custom object
Keras keeps a note of which class generated the config. From the example above, tf.keras.layers.serialize generates a serialized form of the custom layer:
Keras keeps a master list of all built-in layer, model, optimizer, and metric classes, which is used to find the correct class to call from_config . If the class can’t be found, then an error is raised ( Value Error: Unknown layer ). There are a few ways to register custom classes to this list:
- Setting custom_objects argument in the loading function. (see the example in section above «Defining the config methods») or tf.keras.utils.CustomObjectScope
Custom layer and function example
In-memory model cloning
You can also do in-memory cloning of a model via tf.keras.models.clone_model() . This is equivalent to getting the config then recreating the model from its config (so it does not preserve compilation information or layer weights values).
Saving & loading only the model’s weights values
You can choose to only save & load a model’s weights. This can be useful if:
- You only need the model for inference: in this case you won’t need to restart training, so you don’t need the compilation information or optimizer state.
- You are doing transfer learning: in this case you will be training a new model reusing the state of a prior model, so you don’t need the compilation information of the prior model.
APIs for in-memory weight transfer
Weights can be copied between different objects by using get_weights and set_weights :
- tf.keras.layers.Layer.get_weights() : Returns a list of numpy arrays.
- tf.keras.layers.Layer.set_weights() : Sets the model weights to the values in the weights argument.
Transfering weights from one layer to another, in memory
Transfering weights from one model to another model with a compatible architecture, in memory
The case of stateless layers
Because stateless layers do not change the order or number of weights, models can have compatible architectures even if there are extra/missing stateless layers.
APIs for saving weights to disk & loading them back
Weights can be saved to disk by calling model.save_weights in the following formats:
- TensorFlow Checkpoint
The default format for model.save_weights is TensorFlow checkpoint. There are two ways to specify the save format:
- save_format argument: Set the value to save_format=»tf» or save_format=»h5″ .
- path argument: If the path ends with .h5 or .hdf5 , then the HDF5 format is used. Other suffixes will result in a TensorFlow checkpoint unless save_format is set.
There is also an option of retrieving weights as in-memory numpy arrays. Each API has its pros and cons which are detailed below.
TF Checkpoint format
The TensorFlow Checkpoint format saves and restores the weights using object attribute names. For instance, consider the tf.keras.layers.Dense layer. The layer contains two weights: dense.kernel and dense.bias . When the layer is saved to the tf format, the resulting checkpoint contains the keys «kernel» and «bias» and their corresponding weight values. For more information see «Loading mechanics» in the TF Checkpoint guide.
Note that attribute/graph edge is named after the name used in parent object, not the name of the variable. Consider the CustomLayer in the example below. The variable CustomLayer.var is saved with «var» as part of key, not «var_a» .
Transfer learning example
Essentially, as long as two models have the same architecture, they are able to share the same checkpoint.
It is generally recommended to stick to the same API for building models. If you switch between Sequential and Functional, or Functional and subclassed, etc., then always rebuild the pre-trained model and load the pre-trained weights to that model.
Keras to_json(), what does it save?
My impression is that it only saves the model’s architecture, so I should be able to call it before I start training? And then save_weights() saves the weights I need to restore the model? Any more details on this?
At what stage can I call to_json() ? I.e. do I have to call compile() first? Can it be before fit() ?
1 Answer 1
As mentioned in Keras docs it only saves the architecture of the model:
Saving/loading only a model’s architecture
If you only need to save the architecture of a model, and not its weights or its training configuration, you can do:
The generated JSON / YAML files are human-readable and can be manually edited if needed.
You can then build a fresh model from this data:
Training, saving and loading Artificial Neural Networks in Keras
—> Indian Technical Authorship Contest starts on 1st July 2023. Stay tuned.
Reading time: 30 minutes | Coding time: 15 minutes
One of the main reasons why Deep Learning projects take weeks or even months to complete, is that the size of datasets is extremely large. This makes training of the neural network a very long process.
In this article, we will explore the concept of saving your own trained neural network models for the future, thereby having to train it only once, and loading it later.
All the code demonstrations are done in Keras, a high level API based on Tensorflow.
Taking an example of Diabetes detection
Let’s take an example of a Neural Network built for detecting the presence of diabetes in a patient. The dataset used is the Pima Indian Diabetes dataset.
First, let’s import the essential libraries
Fetching and preparing the data
Next, we’ll just import the dataset using Pandas and create the input feature vector (X) and target variable (Y)
Before we go ahead and train our neural network, let’s first have a look at the input features and target variable:
Defining the model structure
Next, we’ll define our simple neural network. The architecture of the neural network has:
- Input layer (8 input features)
- Hidden layer 1 (12 neurons)
- Hidden layer 2 (8 neurons)
- Output layer (1 neuron, softmax function having "0" or "1" as output)
Building the model
Next, we’ll build the model using Adam optimizer:
Training the model
Next, we train the model for 10 epochs
Evaluating the model
Finally, let’s evaluate our model:
Now that we have defined, built, trained and evaluate our model, we may feel the need to save it for later, to avoid going through the entire process again.
Saving the model for later
Saving a model requires 2 steps: Saving the model and saving the weights. The model can be saved to a json file using the to_json() function. The weights are saved to a .h5 file using the save_weights() method. Here’s a code snippet showing the steps:
Once these steps are run, the .json and .h5 files will be created in the local directory. These can be used to load the model as it is in the future. These files are the key for reusing the model.
Loading the model
Now that the model has been saved, let’s try to load the model again and check for accuracy. This is shown in the code snippet below:
It is important to note that you must always compile a model again after loading it, in order for it to run.
Accuracy check on loaded model
Here, we will try and achieve the same accuracy score on our loaded model as we got earlier, i.e. 80.86%.
Hence, we can verify that the performance of our loaded model is exactly same as the model we spent time training. This shows that for other models as well, we can conveniently save a model after training on a large dataset and come back to it later to resume to the same state as earlier.
What we have done, saving a model for later is not limited to one’s own models. Keras supports Transfer learning, allowing us to access pre-trained models by others, thus greatly reducing the amount of time and effort we need to spend.
What is Transfer learning?
Traditionally, in order to train a model with good enough accuracy, we needed to have a large size of the dataset, enough time (a few weeks), and sufficient computational power. This is no longer the case with the approach of Transfer learning. Once a model is trained by someone, the pre-trained model can be used by anyone else for solving their problem using this approach.
Transfer learning basically allows us to use either some or all layers of a neural network originally trained for some other purposes.
Keras allows free and open access to a variety of pre trained models, which anyone can use for training their own modified models with minimal computational effort.
Transfer learning will be the next driver of ML success — Andrew Ng
How to save and load a model with Keras
In this article, we’ll look at how we can save a model that we have trained in Keras.
First, we’ll create a very simple CNN model using Keras and train that model on the Fashion MNIST data set. Then I will show you how to save and subsequently restore that trained model and perform the prediction.
I understand that saving a model for a very basic problem like fashion MNIST might sound trivial.
However, think of a situation where we are training a very complex model which takes a tremendous amount of time, like hours, or even days, to train.
In such cases, it is often useful to save the model. So, in the future, if we need to make predictions, we can simply load the trained model and make our predictions instead of training our model again from scratch.
In Keras, models are saved in the following three formats.
Let’s see how to save the model in each of these formats.
Generally, the models are saved in this format. HDF5 is short for the Hierarchical Data Format. The number five indicates the version.
The HDF format will store our entire model. It saves all the information about our model, including the architecture, model weights, trained parameters, optimizer details etcetera.
In Keras, saving a model in the HDF5 format is quite simple. Keras has a save() function that allows us to easily save the model in this format.