Notebook 07: Updates to Load and Evaluate Checkpoint Weights

[mrdbourke]

After some troubleshooting in the section "Load and evaluate checkpoint weights" of Notebook 07: 07 Milestone Project 1: 🍔👁 Food Vision Big™

As seen in #449, I've found out there may be something wrong with the tf.keras.clone_model method.

I'm not 100% sure on what it could be or whether it's an error with tf.keras.applications.efficientnet (which is notorious for errors across TensorFlow versions).

To fix this, I've decided to keep the topic of the section the same but change the steps.

Only the first two steps are different (and arguably simpler).

Install tf-nightly (to get TensorFlow 2.13.0+)

If there are issues saving the model/using mixed precision with versions of TensorFlow under 2.12.0, you can install tf-nightly (as of May 2023) to get TensorFlow 2.13+.

This version helps to fix issues with saving + loading tf.keras.applications.efficientnet models as well as allows the use of mixed precision training.

You can read more here: https://github.com/keras-team/keras/issues/17199#issuecomment-1501425457

To install tf-nightly in Google Colab, you can use:

!pip install -U tf-nightly

Check the version in the output:

import tensorflow as tf
tf.__version__

>>> 2.14.0-dev20230518 

New loading and evaluating a model via checkpoints

We can load in and evaluate our model's checkpoints by:

1. Recreating a new instance of our model called created_model by turning our original model creation code into a function called create_model().
2. Compiling our created_model with the same loss, optimizer and metrics as the original model (every time you create a new model, you must compile it).
3. Calling the load_weights() method on our cloned model passing it the path to where our checkpointed weights are stored.
4. Calling evaluate() on the cloned model with loaded weights and saving the results.
5. Comparing the created_model results to our previous model results (these should be the exact same, if not very close).

A reminder, checkpoints are helpful for when you perform an experiment such as fine-tuning your model. In the case you fine-tune your feature extraction model and find it doesn't offer any improvements, you can always revert back to the checkpointed version of your model.

# 1. Create a function to recreate the original model
def create_model():
  # Create base model
  input_shape = (224, 224, 3)
  base_model = tf.keras.applications.efficientnet.EfficientNetB0(include_top=False)
  base_model.trainable = False # freeze base model layers

  # Create Functional model 
  inputs = layers.Input(shape=input_shape, name="input_layer")
  # Note: EfficientNetBX models have rescaling built-in but if your model didn't you could have a layer like below
  # x = layers.Rescaling(1./255)(x)
  x = base_model(inputs, training=False) # set base_model to inference mode only
  x = layers.GlobalAveragePooling2D(name="pooling_layer")(x)
  x = layers.Dense(len(class_names))(x) # want one output neuron per class 
  # Separate activation of output layer so we can output float32 activations
  outputs = layers.Activation("softmax", dtype=tf.float32, name="softmax_float32")(x) 
  model = tf.keras.Model(inputs, outputs)
  
  return model

# 2. Create and compile a new version of the original model (new weights)
created_model = create_model()
created_model.compile(loss="sparse_categorical_crossentropy",
                      optimizer=tf.keras.optimizers.Adam(),
                      metrics=["accuracy"])

# 3. Load the saved weights
created_model.load_weights(checkpoint_path)

# 4. Evaluate the model with loaded weights
results_created_model_with_loaded_weights = created_model.evaluate(test_data)

# 5. Compare results with original model
import numpy as np
assert np.isclose(results_feature_extract_model, results_created_model_with_loaded_weights).all(), "Loaded weights results are not close to original model."  # check if all elements in array are close

This code should pass the assertion error.

Old (will error in TensorFlow versions 2.10+)

We can load in and evaluate our model's checkpoints by:

1. Cloning our model using tf.keras.models.clone_model() to make a copy of our feature extraction model with reset weights.
2. Calling the load_weights() method on our cloned model passing it the path to where our checkpointed weights are stored.
3. Calling evaluate() on the cloned model with loaded weights.

A reminder, checkpoints are helpful for when you perform an experiment such as fine-tuning your model. In the case you fine-tune your feature extraction model and find it doesn't offer any improvements, you can always revert back to the checkpointed version of your model.

# Clone the model we created (this resets all weights)
cloned_model = tf.keras.models.clone_model(model)
cloned_model.summary()

# Load checkpointed weights into cloned_model
cloned_model.load_weights(checkpoint_path)

# Compile cloned_model (with same parameters as original model)
cloned_model.compile(loss="sparse_categorical_crossentropy",
                     optimizer=tf.keras.optimizers.Adam(),
                     metrics=["accuracy"])

# Evalaute cloned model with loaded weights (should be same score as trained model)
results_cloned_model_with_loaded_weights = cloned_model.evaluate(test_data)

# Loaded checkpoint weights should return very similar results to checkpoint weights prior to saving
import numpy as np
assert np.isclose(results_feature_extract_model, results_cloned_model_with_loaded_weights).all(), "Loaded weights results are not close to original model."  # check if all elements in array are close