Saving and Loading TensorFlow Models

Github Link

https://github.com/Natan-Asrat/tensorflow_saving_and_loading_models

Contact

The Setup

Introduction

In this project, i used CIFAR-10 dataset with a simple custom CNN model to train and save it using checkpoints.

Model Summary:

Checkpoint including model architecture:

Libraries Used

TensorFlow Keras

Imports

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPooling2D
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.models import load_model

Dataset

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train = x_train / 255.0
x_test = x_test / 255.0

# Use smaller subset -- speeds things up
x_train = x_train[:10000]
y_train = y_train[:10000]
x_test = x_test[:1000]
y_test = y_test[:1000]

Test Accuracy Function

def get_test_accuracy(model, x_test, y_test):
    test_loss, test_acc = model.evaluate(x=x_test, y=y_test, verbose=0)
    print('accuracy: {acc:0.3f}'.format(acc=test_acc))

Function to Create a New Instance of a Simple CNN

def get_new_model():
    model = Sequential([
        Conv2D(filters=16, input_shape=(32, 32, 3), kernel_size=(3, 3), 
               activation='relu', name='conv_1'),
        Conv2D(filters=8, kernel_size=(3, 3), activation='relu', name='conv_2'),
        MaxPooling2D(pool_size=(4, 4), name='pool_1'),
        Flatten(name='flatten'),
        Dense(units=32, activation='relu', name='dense_1'),
        Dense(units=10, activation='softmax', name='dense_2')
    ])
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model

The Code

Create Model

model = get_new_model()

Fit the Model with Checkpoints

filepath='checkpoints/checkpoint'
checkpoint = ModelCheckpoint(filepath=filepath, frequency='epoch', save_weights_only=True, verbose=1)

model.fit(x_train, y_train, epochs=10, validation_split=0.1, callbacks=[checkpoint])

Evaluate the Accuracy of the Model

get_test_accuracy(model, x_test, y_test)

Create a New Model to Load Weights and Evaluate

model = get_new_model()
model.load_weights(filepath)

get_test_accuracy(model, x_test, y_test)

Model Saving Criteria

Create a More Customised Checkpoint with Epoch and Batch

Create TensorFlow checkpoint object with epoch and batch details:

checkpoint_500_path = "checkpoints_500/checkpoint-{epoch:02d}-{batch:04d}" 
checkpoint_500 = ModelCheckpoint(filepath=checkpoint_500_path, save_freq=500, save_weights_only=True, verbose=False)

Checkpoint for Best Accuracy

checkpoint_best_path = "checkpoints_best/checkpoint"
checkpoint_best = ModelCheckpoint(filepath=checkpoint_best_path,
                                 save_freq='epoch',
                                 save_weights_only=True,
                                 monitor='val_accuracy',
                                 save_best_only=True,
                                 verbose=1)

Saving the Entire Model

Create checkpoint that saves whole model, not just weights:

checkpoint_entire_model_path = "checkpoints_entire_model"
checkpoint_entire_model = ModelCheckpoint(
    filepath=checkpoint_entire_model_path,
    verbose=1,
    save_weights_only=False,
    frequency='epoch'
)

Loading the model:

model = load_model(checkpoint_entire_model_path)

Natville