TensorFlow is an open-source machine-learning library developed by Google. In this article, we are going to explore the how can we load a model in TensorFlow.

tf.keras.models.load_model

tf.keras.models.load_model function is used to load saved models from storage for further use. It allows users to easily retrieve trained models from disk or other storage mediums.

The syntax of the tf.keras.models.load_model function is as follows:

tf.keras.models.load_model(filepath, custom_objects=None, compile=True)

where,

• file path: This argument specifies the path to the saved model file or an h5py.File object from which to load the model.
• custom_objects: An optional dictionary that maps names (strings) to custom classes or functions to be considered during deserialization. This parameter proves invaluable when dealing with models containing custom layers or loss functions.
• compile: A boolean flag indicating whether to compile the loaded model after loading. When set to True, the model will be compiled, leveraging the optimizer and loss function specified during training. Conversely, setting it to False allows for skipping compilation, useful when solely interested in model inference.

Implementation of tf.keras.models.load_model in TensorFlow

Importing Necessary Libraries

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

Define a convolutional neural network model

This code defines a simple CNN model with three convolutional layers followed by max pooling, flattening, and two dense layers for classification. The model takes input images of size 128×128 pixels with 3 channels (RGB) and outputs a probability distribution over 10 classes using SoftMax activation.

# Create a simple CNN model
model = Sequential([
    Conv2D(32, (3, 3), activation='relu', input_shape=(128, 128, 3)),
    MaxPooling2D((2, 2)),
    Conv2D(64, (3, 3), activation='relu'),
    MaxPooling2D((2, 2)),
    Conv2D(64, (3, 3), activation='relu'),
    Flatten(),
    Dense(64, activation='relu'),
    Dense(10, activation='softmax')
])

Compile and Save the Model

# Compile the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Save the model
model.save('model.h5')

Loading The Model

We load a saved model from the file ‘model.h5’ using TensorFlow‘s load_model function and then prints a summary of the loaded model, showing the model architecture, layer names, output shapes, and number of parameters.

# Load the saved model
loaded_model = tf.keras.models.load_model('model.h5')

loaded_model.summary()

Output:

Model: "sequential_3"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 conv2d_6 (Conv2D)           (None, 126, 126, 32)      896       
                                                                 
 max_pooling2d_4 (MaxPoolin  (None, 63, 63, 32)        0         
 g2D)                                                            
                                                                 
 conv2d_7 (Conv2D)           (None, 61, 61, 64)        18496     
                                                                 
 max_pooling2d_5 (MaxPoolin  (None, 30, 30, 64)        0         
 g2D)                                                            
                                                                 
 conv2d_8 (Conv2D)           (None, 28, 28, 64)        36928     
                                                                 
 flatten_2 (Flatten)         (None, 50176)             0         
                                                                 
 dense_5 (Dense)             (None, 64)                3211328   
                                                                 
 dense_6 (Dense)             (None, 10)                650       
                                                                 
=================================================================
Total params: 3268298 (12.47 MB)
Trainable params: 3268298 (12.47 MB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________

Conclusion

In conclusion, the tf.keras.models.load_model function is a powerful tool for loading saved Keras models in TensorFlow. By understanding its usage and arguments, developers can seamlessly integrate saved models into their applications, enabling efficient model deployment and inference.