Saving and serializing 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)

# 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)

# 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)

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

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)

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

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

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

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)

# 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)

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

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.