Notebook 05: Updates to original notebook + fix TypeError: Unable to serialize

[mrdbourke]

The following note keeps track of some relatively minor changes in Notebook 05: Transfer Learning with TensorFlow Part 2: Fine-tuning to migitate a few issues.

The main one being: tf.keras.applications.efficientnet being notoriously buggy across TensorFlow versions.

The following code fixes should fix all issues in Notebook 05 and has been test on TensorFlow versions 2.12.0 and 2.13.0

TODO

• add a note to ZTM lectures
• add a note to Udemy lectures
• add a note to the README about the update
• add a note about using EffNetB0 (V1) vs EffNetB0V2 (perhaps less errors with V2)
• go through notebook 05 and update writing for EffNetB0V2
• every line that used effnetb0 is now effnetB0v2

Quick summary

• Swapped all references of EfficientNetB0 -> EfficientNetV2B0 to fix Notebook 05: TypeError: Unable to serialize [2.0897 2.1129 2.1082] to JSON. Unrecognized type <class 'tensorflow.python.framework.ops.EagerTensor'> (fix inside) #553 (if you are getting errors with EfficientNetB0, try using EfficientNetV2B0)
• Functionized model_2 creation to fix Notebook 05: load_weights results in "Incompatible tensor with shape (1280, 10)..." #544
• Made referencing the base_model clearer

Quick links:

• New notebook (updated for August 2023), tested for TensorFlow 2.12.0 and 2.13.0: https://github.com/mrdbourke/tensorflow-deep-learning/blob/main/05_transfer_learning_in_tensorflow_part_2_fine_tuning.ipynb
• Old notebook (for reference): https://colab.research.google.com/drive/1rerx_h0fNMCgu-j1y15a0dZ_xDKBINrc?usp=sharing

EfficienetNetB0 -> EfficientNetV2B0

In versions of TensorFlow 2.10+, the tf.keras.applications.efficientnet.EfficientNetB0 (EfficientNetB0 for short) module seemed to be plagued with errors.

To fix this, all uses of EfficientNetB0 have been replaced with tf.keras.applications.efficientnet_v2.EfficientNetV2B0 (EfficientNetV2B0 for short).

New:

base_model = tf.keras.applications.efficientnet_v2.EfficientNetV2B0(include_top=False)

Old:

base_model = tf.keras.applications.efficientnet.EfficientNetB0(include_top=False)

Why?

Mostly because I've found less bugs (alongside other students of the course, TK - Moophers GitHub).

Note: This doesn't necessarily mean EfficientNetV2B0 is the best model for the job, it performs well on many computer vision tasks, however, it's only used as an example. Best to experiment with other model versions in tf.keras.applications to find the best model for your own problem.

This takes care of the issue in #553.

Functionizing model_2 creation

Because model_2 is used several times throughout Notebook 05, I made a reference function to recreate it.

That way we know whenever the function is called, we're getting a new instance of model_2 (rather than potentially reusing an old one).

def create_base_model(input_shape: tuple[int, int, int] = (224, 224, 3), 
                      output_shape: int = 10, 
                      learning_rate: float = 0.001,
                      training: bool = False) -> tf.keras.Model:
    """
    Create a model based on EfficientNetV2B0 with built-in data augmentation.

    Parameters:
    - input_shape (tuple): Expected shape of input images. Default is (224, 224, 3).
    - output_shape (int): Number of classes for the output layer. Default is 10.
    - learning_rate (float): Learning rate for the Adam optimizer. Default is 0.001.
    - training (bool): Whether the base model is trainable. Default is False.

    Returns:
    - tf.keras.Model: The compiled model with specified input and output settings.
    """
  
    # Create base model
    base_model = tf.keras.applications.efficientnet_v2.EfficientNetV2B0(include_top=False)
    base_model.trainable = training

    # Setup model input and outputs with data augmentation built-in
    inputs = layers.Input(shape=input_shape, name="input_layer") 
    x = data_augmentation(inputs) 
    x = base_model(x, training=False)  # pass augmented images to base model but keep it in inference mode
    x = layers.GlobalAveragePooling2D(name="global_average_pooling_layer")(x)
    outputs = layers.Dense(units=output_shape, activation="softmax", name="output_layer")(x)
    model = tf.keras.Model(inputs, outputs)

    # Compile model
    model.compile(loss="categorical_crossentropy",
                  optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
                  metrics=["accuracy"])
  
    return model

# Create an instance of model_2 with our new function
model_2 = create_base_model()

This way, for experiment 3 and experiment 4, we can use a new instance of model_2, load the weights from experiment 2 and then fine-tune appropriately.

This helps to fix the issue in #544 (though this issue is also related to EfficientNetB0 being notoriously buggy across TensorFlow versions).

Change base_model -> model_2_base_model

Using the same base_model variable throughout the notebook got confusing.

So I updated it to be model_2_base_model so we know which base it relates to.

For example:

model_2 = create_base_model()

# Unfreeze the top 10 layers in model_2's base_model
model_2_base_model = model_2.layers[2]
model_2_base_model.trainable = True

# Freeze all layers except for the last 10
for layer in model_2_base_model.layers[:-10]:
  layer.trainable = False