Saving and serializing models

The first part of this guide covers saving and serialization for Sequential models and models built using the Functional API. The saving and serialization APIs are the exact same for both of these types of models.

Saving for custom subclasses of Model is covered in the section “Saving Subclassed Models”. The APIs in this case are slightly different than for Sequential or Functional models.

Overview

For Sequential Models and models built using the Functional API use:

save_model_hdf5()/load_model_hdf5() to save the entire model to disk, including the optimizer state. You can also use save_model_tf/load_model_tf to save the entire model to the SavedModel format.
get_config()/from_config() to load only the model architecture into an R object.
model_to_json()/model_from_json() to save only the architecture of the model to a single string - useful for saving the architecture to disk. You can also use model_to_yaml()/model_from_yaml() to save the architecture.
save_model_weights_hdf5()/load_model_weights_hdf5() if you want to save only the model weights to disk in the hdf5 format. You can also use save_model_weights_tf()/load_model_weights_tf() to save the weights in the SavedModel format.

Note you can use a combination of model_to_json() and save_model_weights_hdf5() to save both the architecture and the weights. In this case the optimizer state is not saved.

For custom models use:

save_model_weights_tf() or save_model_weights_hdf5() to save the model weights. Usually for custom models, the architecture must be recreated using code.

Setup

library(keras)

Saving Sequential Models or Functional models

inputs <- layer_input(shape = 784, name = "digits")
outputs <- inputs %>% 
  layer_dense(units = 64, activation = "relu", name = "dense_1") %>% 
  layer_dense(units = 64, activation = "relu", name = "dense_2") %>% 
  layer_dense(units = 10, activation = "softmax", name = "predictions")
model <- keras_model(inputs, outputs) 
summary(model)

Optionally, let’s train this model, just so it has weight values to save, as well as an an optimizer state. Of course, you can save models you’ve never trained, too, but obviously that’s less interesting.

c(c(x_train, y_train), c(x_test, y_test)) %<-% dataset_mnist()
x_train <- x_train %>% array_reshape(dim = c(60000, 784))/255
x_test <- x_test %>% array_reshape(dim = c(10000, 784))/255

model %>% compile(loss = "sparse_categorical_crossentropy",
                  optimizer = optimizer_rmsprop())

history <- model %>% fit(x_train, y_train, batch_size = 64, epochs = 1)

# Save predictions for future checks
predictions <- predict(model, x_test)

Whole-model saving

You can save a model built with the Functional API into a single file. You can later recreate the same model from this file, even if you no longer have access to the code that created the model.

This file includes:

The model’s architecture
The model’s weight values (which were learned during training)
The model’s training config (what you passed to compile), if any
The optimizer and its state, if any (this enables you to restart training where you left off)

# Save the model
save_model_hdf5(model, "model.h5")

# Recreate the exact same model purely from the file
new_model <- load_model_hdf5("model.h5")

# Check that the state is preserved
new_predictions <- predict(new_model, x_test)
all.equal(predictions, new_predictions)

Note that the optimizer state is preserved as well so you can resume training where you left off.

Export to SavedModel

You can also export a whole model to the TensorFlow SavedModel format. SavedModel is a standalone serialization format for Tensorflow objects, supported by TensorFlow serving as well as TensorFlow implementations other than Python. Note that save_model_tf is only available for TensorFlow version greater than 1.14.

# Export the model to a SavedModel
save_model_tf(model, "model/")

# Recreate the exact same model
new_model <- load_model_tf("model/")

# Check that the state is preserved
new_predictions <- predict(new_model, x_test)
all.equal(predictions, new_predictions)

Note that the optimizer state is preserved as well so you can resume training where you left off.

The SavedModel files that were created contain:

A TensorFlow checkpoint containing the model weights.
A SavedModel proto containing the underlying Tensorflow graph. Separate graphs are saved for prediction (serving), train, and evaluation. If the model wasn’t compiled before, then only the inference graph gets exported.
The model’s architecture config, if available.

You can also use the export_savedmodel function to export models but those models can not be loaded as Keras models again. Models exported using exported_savedmodels can be used for prediction though.

export_savedmodel(model, "savedmodel/")
new_predictions <- tfdeploy::predict_savedmodel(x_test, "savedmodel/")

Note Exporting with export_savedmodel sets learning phase to 0 so you need to restart R and re-build the model before doing additional training.

Architecture-only saving

Sometimes, you are only interested in the architecture of the model, and you don’t need to save the weight values or the optimizer. In this case, you can retrieve the “config” of the model via the get_config() method. The config is a named list that enables you to recreate the same model – initialized from scratch, without any of the information learned previously during training.

config <- get_config(model)
reinitialized_model <- from_config(config)

# Note that the model state is not preserved! We only saved the architecture.
new_predictions <- predict(reinitialized_model, x_test)
all.equal(predictions, new_predictions)

You can alternatively use model_to_json() and model_from_json(), which uses a JSON string to store the config instead of a named list. This is useful to save the config to disk.

json_config <- model_to_json(model)
reinitialized_model <- model_from_json(json_config)

Weights-only saving

Sometimes, you are only interested in the state of the model – its weights values – and not in the architecture. In this case, you can retrieve the weights values as a list of arrays via get_weights(), and set the state of the model via set_weights:

weights <- get_weights(model)
set_weights(reinitialized_model, weights)

new_predictions <- predict(reinitialized_model, x_test)
all.equal(predictions, new_predictions)

You can combine get_config()/from_config() and get_weights()/set_weights() to recreate your model in the same state. However, unlike save_model_hdf5, this will not include the training config and the optimizer. You would have to call compile() again before using the model for training.

config <- get_config(model)
weights <- get_weights(model)

new_model <- from_config(config)
set_weights(new_model, weights)

# Check that the state is preserved
new_predictions <- predict(new_model, x_test)
all.equal(predictions, new_predictions)

Note that the optimizer was not preserved, so the model should be compiled anew before training (and the optimizer will start from a blank state).

The save-to-disk alternative to get_weights() and set_weights(weights) is save_weights(fpath) and load_weights(fpath).

# Save JSON config to disk
json_config <- model_to_json(model)
writeLines(json_config, "model_config.json")

# Save weights to disk
save_model_weights_hdf5(model, "model_weights.h5")

# Reload the model from the 2 files we saved
json_config <- readLines("model_config.json")
new_model <- model_from_json(json_config)
load_model_weights_hdf5(new_model, "model_weights.h5")

# Check that the state is preserved
new_predictions <- predict(new_model, x_test)
all.equal(predictions, new_predictions)

Note that the optimizer was not preserved. But remember that the simplest, recommended way is just this:

save_model_hdf5(model, "model.h5")
new_model <- load_model_hdf5("model.h5")

Weights-only saving in SavedModel format

Note that save_weights can create files either in the Keras HDF5 format, or in the TensorFlow SavedModel format.

save_model_weights_tf(model, "model_weights")

Saving Subclassed Models

Sequential models and Functional models are data structures that represent a DAG of layers. As such, they can be safely serialized and deserialized.

A subclassed model differs in that it’s not a data structure, it’s a piece of code. The architecture of the model is defined via the body of the call method. This means that the architecture of the model cannot be safely serialized. To load a model, you’ll need to have access to the code that created it (the code of the model subclass). Alternatively, you could be serializing this code as bytecode (e.g. via pickling), but that’s unsafe and generally not portable.

For more information about these differences, see the article “What are Symbolic and Imperative APIs in TensorFlow 2.0?”.

Let’s consider the following subclassed model, which follows the same structure as the model from the first section:

keras_model_simple_mlp <- function(num_classes, 
                                   use_bn = FALSE, use_dp = FALSE, 
                                   name = NULL) {
  
  # define and return a custom model
  keras_model_custom(name = name, function(self) {
    
    # create layers we'll need for the call (this code executes once)
    self$dense1 <- layer_dense(units = 32, activation = "relu")
    self$dense2 <- layer_dense(units = num_classes, activation = "softmax")
    if (use_dp)
      self$dp <- layer_dropout(rate = 0.5)
    if (use_bn)
      self$bn <- layer_batch_normalization(axis = -1)
    
    # implement call (this code executes during training & inference)
    function(inputs, mask = NULL) {
      x <- self$dense1(inputs)
      if (use_dp)
        x <- self$dp(x)
      if (use_bn)
        x <- self$bn(x)
      self$dense2(x)
    }
  })
}

model <- keras_model_simple_mlp(num_classes = 10)

First of all, a subclassed model that has never been used cannot be saved.

That’s because a subclassed model needs to be called on some data in order to create its weights.

Until the model has been called, it does not know the shape and dtype of the input data it should be expecting, and thus cannot create its weight variables. You may remember that in the Functional model from the first section, the shape and dtype of the inputs was specified in advance (via layer_input) – that’s why Functional models have a state as soon as they’re instantiated.

Let’s train the model, so as to give it a state:

model %>% compile(loss = "sparse_categorical_crossentropy",
                  optimizer = optimizer_rmsprop())

history <- model %>% fit(x_train, y_train, batch_size = 64, epochs = 1)

The recommended way to save a subclassed model is to use save_model_weights_tf to create a TensorFlow SavedModel checkpoint, which will contain the value of all variables associated with the model: - The layers’ weights - The optimizer’s state - Any variables associated with stateful model metrics (if any).

save_model_weights_tf(model, "my_weights")

# Save predictions for future checks
predictions <- predict(model, x_test)
# Also save the loss on the first batch
# to later assert that the optimizer state was preserved
first_batch_loss <- train_on_batch(model, x_train[1:64,], y_train[1:64])

To restore your model, you will need access to the code that created the model object.

Note that in order to restore the optimizer state and the state of any stateful metric, you should compile the model (with the exact same arguments as before) and call it on some data before calling load_weights:

new_model <- keras_model_simple_mlp(num_classes = 10)
new_model %>% compile(loss = "sparse_categorical_crossentropy",
                  optimizer = optimizer_rmsprop())

# This initializes the variables used by the optimizers,
# as well as any stateful metric variables
train_on_batch(new_model, x_train[1:5,], y_train[1:5])

# Load the state of the old model
load_model_weights_tf(new_model, "my_weights")

# Check that the model state has been preserved
new_predictions <- predict(new_model, x_test)
all.equal(predictions, new_predictions)

# The optimizer state is preserved as well,
# so you can resume training where you left off
new_first_batch_loss <- train_on_batch(new_model, x_train[1:64,], y_train[1:64])
first_batch_loss == new_first_batch_loss

You’ve reached the end of this guide! This covers everything you need to know about saving and serializing models with Keras in TensorFlow 2.0.