SavedModel format in TensorFlow - ML Experiment: Train & Evaluate

Experiment - SavedModel format

Problem:You have trained a TensorFlow model and saved it using the SavedModel format. However, when you reload the model, it does not produce the same predictions as before saving.

Current Metrics:Original model prediction on test input: [0.85, 0.15]; Reloaded model prediction on same input: [0.60, 0.40]

Issue:The reloaded model predictions differ significantly from the original model, indicating a problem with saving or loading the model.

Your Task

Ensure that the TensorFlow model saved in SavedModel format can be reloaded and produce the same predictions as the original model.

Use TensorFlow 2.x APIs only

Do not retrain the model after loading

Use the SavedModel format for saving and loading

Hint 1

Hint 2

Hint 3

Solution

TensorFlow

import tensorflow as tf
import numpy as np

# Create a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(5,)),
    tf.keras.layers.Dense(2, activation='softmax')
])

# Compile and train the model on dummy data
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

X_train = np.random.random((100, 5))
y_train = np.random.randint(0, 2, 100)
model.fit(X_train, y_train, epochs=3, batch_size=10, verbose=0)

# Test input
X_test = np.random.random((1, 5))
original_pred = model.predict(X_test)

# Save the model in SavedModel format
model.save('saved_model_dir')

# Load the model back
loaded_model = tf.keras.models.load_model('saved_model_dir')

# Predict with loaded model
loaded_pred = loaded_model.predict(X_test)

print(f'Original prediction: {original_pred}')
print(f'Loaded model prediction: {loaded_pred}')

Used model.save() with a directory path to save in SavedModel format

Used tf.keras.models.load_model() to reload the model correctly

Verified predictions before and after saving to ensure they match

Results Interpretation

Before saving, the model predicted probabilities like [0.85, 0.15]. After saving and loading, the predictions changed to [0.60, 0.40], which is incorrect.

After applying the correct saving and loading methods, both predictions match closely, for example, [0.72, 0.28] before and after.

Using TensorFlow's SavedModel format with proper save and load functions preserves the model's architecture and weights, ensuring consistent predictions before and after saving.

Bonus Experiment

Try saving and loading a model that includes a custom layer or custom loss function using the SavedModel format.

💡 Hint

You will need to provide a custom_objects dictionary when loading the model to tell TensorFlow how to handle the custom parts.