Save and load TensorFlow model

Page content

Original page

https://www.tensorflow.org/tutorials/keras/save_and_load

Show model details

import tensorflow as tf

def create_model():
  model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
  ])

  model.compile(optimizer='adam',
                loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['accuracy'])

  return model

# Create a basic model instance
model = create_model()

# Display the model's architecture
model.summary()

Load MNIST data

# MNIST
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

train_labels = train_labels[:1000]
test_labels = test_labels[:1000]

train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0

Save model while training (checkpoint)

Using tf.keras.callbacks.

import os
checkpoint_path = "save_epoch/cp.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

# Create a callback that saves the model's weights
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)

# Train the model with the new callback
model.fit(train_images, 
          train_labels,  
          epochs=10,
          validation_data=(test_images,test_labels),
          callbacks=[cp_callback])  # Pass callback to training

# This may generate warnings related to saving the state of the optimizer.
# These warnings (and similar warnings throughout this notebook)
# are in place to discourage outdated usage, and can be ignored.

Here is a full output in my case.

Epoch 1/10
32/32 [==============================] - ETA: 0s - loss: 1.2029 - accuracy: 0.6510
Epoch 00001: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 10ms/step - loss: 1.2029 - accuracy: 0.6510 - val_loss: 0.7293 - val_accuracy: 0.7840
Epoch 2/10
32/32 [==============================] - ETA: 0s - loss: 0.4511 - accuracy: 0.8600
Epoch 00002: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 5ms/step - loss: 0.4511 - accuracy: 0.8600 - val_loss: 0.5207 - val_accuracy: 0.8360
Epoch 3/10
32/32 [==============================] - ETA: 0s - loss: 0.2969 - accuracy: 0.9210
Epoch 00003: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 5ms/step - loss: 0.2969 - accuracy: 0.9210 - val_loss: 0.4866 - val_accuracy: 0.8430
Epoch 4/10
32/32 [==============================] - ETA: 0s - loss: 0.2145 - accuracy: 0.9530
Epoch 00004: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 5ms/step - loss: 0.2145 - accuracy: 0.9530 - val_loss: 0.4434 - val_accuracy: 0.8530
Epoch 5/10
 1/32 [..............................] - ETA: 0s - loss: 0.1201 - accuracy: 1.0000
Epoch 00005: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 4ms/step - loss: 0.1531 - accuracy: 0.9680 - val_loss: 0.4229 - val_accuracy: 0.8650
Epoch 6/10
 1/32 [..............................] - ETA: 0s - loss: 0.2162 - accuracy: 0.9062
Epoch 00006: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 4ms/step - loss: 0.1150 - accuracy: 0.9790 - val_loss: 0.3989 - val_accuracy: 0.8690
Epoch 7/10
 1/32 [..............................] - ETA: 0s - loss: 0.0774 - accuracy: 1.0000
Epoch 00007: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 4ms/step - loss: 0.0907 - accuracy: 0.9830 - val_loss: 0.4240 - val_accuracy: 0.8550
Epoch 8/10
 1/32 [..............................] - ETA: 0s - loss: 0.0437 - accuracy: 1.0000
Epoch 00008: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 4ms/step - loss: 0.0664 - accuracy: 0.9920 - val_loss: 0.3999 - val_accuracy: 0.8620
Epoch 9/10
 1/32 [..............................] - ETA: 0s - loss: 0.0272 - accuracy: 1.0000
Epoch 00009: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 5ms/step - loss: 0.0631 - accuracy: 0.9940 - val_loss: 0.3949 - val_accuracy: 0.8690
Epoch 10/10
 1/32 [..............................] - ETA: 0s - loss: 0.0240 - accuracy: 1.0000
Epoch 00010: saving model to training_1/cp.ckpt
32/32 [==============================] - 0s 4ms/step - loss: 0.0447 - accuracy: 0.9990 - val_loss: 0.4180 - val_accuracy: 0.8660

It returns models of each checkpoints.

$ ls training_1/
checkpoint  cp.ckpt.data-00000-of-00002  cp.ckpt.data-00001-of-00002  cp.ckpt.index

checkpoint file just describe a data paths.

$ cat training_1/checkpoint
model_checkpoint_path: "cp.ckpt"
all_model_checkpoint_paths: "cp.ckpt"

Other files indicate weights of saved model of each check points. The data format is called HDF5.

$ cat training_1/cp.ckpt.index

_CHECKPOINTABLE_OBJECT_GRAP(�5K��4layer_with_weights-0/bias/.ATTRIBUTES/VARIABLE_VAL� ��b(�5�
"OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VAL� ���(�5v��4v/.ATTRIBUTES/VARIABLE_VAL� ��(�5��_!kernel/.ATTRIBUTES/VARIABLE_VALUE
�(��b5X�ǘ5OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE
� ��c(��b5[s�,6v/.ATTRIBUTES/VARIABLE_VALUE
� ���(��b5��!1/bias/.ATTRIBUTES/VARIABLE_VALUE
 ��c((5x��$OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE
 ���((5o�4v/.ATTRIBUTES/VARIABLE_VALUE
 萪((5���:!kernel/.ATTRIBUTES/VARIABLE_VALUE   �
 ��b(��5E�,O5OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE     �
 ���(��5U6v/.ATTRIBUTES/VARIABLE_VALUE �
 ��(��5���+optimizer/beta_1/.ATTRIBUTES/VARIABLE_VALU ��c(5����2/.ATTRIBUTES/VARIABLE_VALU ��c(5����*optimizer/decay/.ATTRIBUTES/VARIABLE_VALU ��c(5�{:
iter/.ATTRIBUTES/VARIABLE_VALUE  ��c5.�dM
(learning_rate/.ATTRIBUTES/VARIABLE_VALU ��c(5i�U2����b��p�	���ɒ  ��	W�$uG

In this example, we save weights, not model it self. If you want to load saved model again, you should create the identical model and load the weights.

How to load

Open other python file in a same directory and write as follows.

import tensorflow as tf

# create the same model
def create_model():
  model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10)
  ])
  model.compile(optimizer='adam',
                loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['accuracy'])
  return model
model = create_model()

# Evaluate first random model
loss, acc = model.evaluate(test_images,  test_labels, verbose=2)

model.load_weights(checkpoint_path)

# Re-evaluate the model
loss,acc = model.evaluate(test_images,  test_labels, verbose=2)
print("Restored model, accuracy: {:5.2f}%".format(100*acc))

Manually save a weight

Caution. It doesn’t store model, weight instead.

# Save the weights
model.save_weights('./checkpoints/my_checkpoint')

Save an entire model

# Create and train a new model instance.
model = create_model()
model.fit(train_images, train_labels, epochs=5)

# Save the entire model as a SavedModel.
!mkdir -p saved_model
model.save('saved_model/my_model') 

# How to load
new_model = tf.keras.models.load_model('saved_model/my_model')