Saving and restoring a TensorFlow model

If we want to use our machine learning model in production or reuse our trained model for a transfer learning task, we have to store our model. In this section, we will outline some methods for storing and restoring the weights or the whole model.

Getting ready

In this recipe, we want to summarize various ways to store a TensorFlow model. We will cover the best way to save and restore an entire model, only the weights, and model checkpoints.

How to do it...

1. We start by loading the necessary libraries:

   import tensorflow as tf
   

2. Next, we'll build an MNIST model using the Keras Sequential API:

   (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
   # Normalize
   x_train = x_train / 255
   x_test = x_test/ 255
   model = tf.keras.Sequential()
   model.add(tf.keras.layers.Flatten(name="FLATTEN"))
   model.add(tf.keras.layers.Dense(units=128 , activation="relu", name...