Saving and Loading a TensorFlow Model

Welcome to this next lesson on Saving and Loading a TensorFlow Model. By the end of this lesson, you’ll be able to understand the importance of saving and loading models, how to save a trained TensorFlow model, load it from the saved file format, and validate the loaded model. This will give you a full cycle understanding and hands-on knowledge on how to handle models when training is done. With the provided code examples that train, save, load, and test a model, let's start our lesson!

When building machine learning models, it's important to save your models for various reasons. The most obvious one is efficiency - once you trained an intricate model that could take hours or even days to train, you want to keep the learned weights to avoid re-training. So, you’d save the model for reuse later without the need to retrain.

Not only that, but the saved model can be shared with others - if you're collaborating with other professionals or even publishing your results, it aids in reproducibility of your results by others. Finally, when deploying a model to production you'll need to load the trained model to make predictions on new data.

In the previous lessons, we trained a TensorFlow model. Now, let's save it!

Before we focus on saving our model, let's briefly revisit the key steps we took to load our data and train the model. Here's the code snippet we used to preprocess the Iris dataset:

import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder

def load_preprocessed_data():
    # Load the Iris dataset
    iris = load_iris()
    X, y = iris.data, iris.target

    # Split the dataset into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, stratify=y, random_state=42)

    # Scale the features
    scaler = StandardScaler().fit(X_train)
    X_train_scaled = scaler.transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # One-hot encode the targets
    encoder = OneHotEncoder(sparse_output=False).fit(y_train.reshape(-1, 1))
    y_train_encoded = encoder.transform(y_train.reshape(-1, 1))
    y_test_encoded = encoder.transform(y_test.reshape(-1, 1))

    return X_train_scaled, X_test_scaled, y_train_encoded, y_test_encoded

And to train a model with our preprocessed data:

import tensorflow as tf
from data_preprocessing import load_preprocessed_data

# Load preprocessed data
X_train, X_test, y_train, y_test = load_preprocessed_data()

# Define the model
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(4,)),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax')
])

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

# Train the model
history = model.fit(X_train, y_train, epochs=150, batch_size=5, validation_data=(X_test, y_test))

In summary:

1. We began by loading the preprocessed data using the load_preprocessed_data() function from our data_preprocessing.py file to get our datasets: X_train, X_test, y_train, and y_test.
2. We constructed a sequential model with TensorFlow's Keras API, with the input shape tailored to our dataset, including several dense layers with ReLU and Softmax activations.
3. The model was then compiled using the Adam optimizer and the categorical crossentropy loss function, with accuracy as a metric.
4. Finally, we trained the model for 150 epochs with a batch size of 5, validating its performance on the test data throughout the training process.

With our training steps revisited, let's move on to saving our well-trained model.