Save and Load Models using TensorFlow in Json?

If you are looking to explore Machine Learning with TensorFlow, you are at the right place. This comprehensive article explains how to save and load the models in TensorFlow along with its brief overview. If you read this article till the end, you will not need to look for further guides on how to save and reuse the Model in Machine Learning.

TensorFlow has become the top-notch choice among Machine Learning Experts. This is because it offers a lot of high-level APIs and pre-built modules to create and train the Machine Learning Models. Thus, it becomes important to learn how to save and load models using the TensorFlow Library. There is not one way to do it, there are various methods. So, let us see what method will be the best one for saving the model object and loading it back from the memory.

What is TensorFlow?

Google Brain Team developed TensorFlow to build and train the Deep Learning Models. It is now open-source, and you can use it to develop your Machine Learning Applications. It provides an efficient set of tools for numerical computation. It follows the approach of the Computation Graph in which the Nodes represent the mathematical operations and edges represent the flow of data between the operations.

It provides various tools and APIs for large-scale machine-learning tasks like image recognition, natural language processing, and reinforcement learning. One of its key features is the Keras, which is High-level API to build and deploy the Machine Learning Models. We will now create and train the model using TensorFlow so that we can save and reuse it.

How to create Model in TensorFlow?

Let us create the sample Model using TensorFlow that classifies the images of the clothing from the MNIST Dataset. The below code shows the model building for this example. First, we have to download and preprocess the Fashion MNIST Data. Then, we have to create and train the neural network using the Sequential API of TensorFlow with the Flatten, Dense, and Dropout Layers.

After building the model, we have to compile it using the Adam optimizer along with the loss function. Then, we predict the output for the sample input which is shown in the output of the code snippet.

# Import TensorFlow and Fashion MNIST dataset
import tensorflow as tf
from tensorflow.keras.datasets import fashion_mnist


# Load and preprocess the Fashion MNIST dataset
(x_train, y_train), (x_test, y_test) = fashion_mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0


# Define the model architecture
model = tf.keras.models.Sequential([
    # Flatten the 28x28 input images into 1D array
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    # Fully connected layer with 128 neurons and ReLU activation
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),  # Dropout layer to prevent overfitting
    # Output layer with 10 neurons for 10 classes and softmax activation
    tf.keras.layers.Dense(10, activation='softmax')
])


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


# Train the model
model.fit(x_train, y_train, epochs=5)


# Evaluate the model on test data
test_loss, test_acc = model.evaluate(x_test, y_test)


print('\nTest accuracy:', test_acc)


# Make predictions on some test data and display predicted labels
predictions = model.predict(x_test[:5])
predicted_labels = [tf.argmax(prediction).numpy()
                    for prediction in predictions]


print('Predicted labels:', predicted_labels)

Output:

Epoch 1/5
1875/1875 [==============================] - 15s 7ms/step - loss: 0.5356 - accuracy: 0.8116
Epoch 2/5
1875/1875 [==============================] - 8s 4ms/step - loss: 0.4016 - accuracy: 0.8550
Epoch 3/5
1875/1875 [==============================] - 10s 6ms/step - loss: 0.3688 - accuracy: 0.8646
Epoch 4/5
1875/1875 [==============================] - 13s 7ms/step - loss: 0.3475 - accuracy: 0.8720
Epoch 5/5
1875/1875 [==============================] - 9s 5ms/step - loss: 0.3302 - accuracy: 0.8798
313/313 [==============================] - 1s 2ms/step - loss: 0.3617 - accuracy: 0.8721
Test accuracy: 0.8720999956130981
1/1 [==============================] - 0s 96ms/step
Predicted labels: [9, 2, 1, 1, 6]

Save and Load Model in TensorFlow

In this method, TensorFlow saves only the model architecture. To do this, it serializes the model architecture into JSON String which contains all the configuration details like layers and parameters. And when we call the load() method, TensorFlow uses this JSON String to reconstruct the model. Following code demonstrates this:

# Save the model architecture to JSON file
model_json = model.to_json()
with open('my_model.json', 'w') as json_file:
    json_file.write(model_json)


# Output confirmation message
print("Model architecture saved successfully.")


# Load the model architecture from JSON file
with open('my_model.json', 'r') as json_file:
    loaded_model_json = json_file.read()
loaded_model = tf.keras.models.model_from_json(loaded_model_json)


# Output confirmation message
print("Model architecture loaded successfully.")

Output:

Model architecture saved successfully.
Model architecture loaded successfully.

The json file gets saved as "my_model.json".

Conclusion

TensorFlow provides various tools, libraries, APIs, and modules for building and saving Machine Learning Models. Thus, we can easily preserve the model’s architecture and reuse it when required.

Now, you can easily save and load the model in TensorFlow.