Saving and Loading Models

Once you've invested time and computational resources into training a deep learning model, you certainly don't want to lose that progress or have to retrain it every time you need to use it. Saving your trained model allows you to pause and resume training, share your model with others, deploy it into applications, or compare different versions. Standard methods for saving and loading Keras models and their learned parameters (weights) are outlined.

Why Save Your Models?

Saving models is a fundamental part of the machine learning workflow for several reasons:

1. Persistence: Training can take hours or even days. Saving allows you to preserve the trained state without needing to retrain from scratch.
2. Deployment: Saved models can be loaded into different environments or applications to make predictions on new data.
3. Sharing and Collaboration: You can share your trained models with colleagues or the community.
4. Fine-tuning and Transfer Learning: You can load a pre-trained model and continue training it on a new dataset or task.
5. Experimentation: Save different versions of a model (e.g., trained with different hyperparameters) for later comparison.
6. Fault Tolerance: Combined with callbacks like ModelCheckpoint (discussed previously), saving allows you to recover the best performing model state even if training is interrupted.

Saving the Entire Model

The most comprehensive way to save your work is to save the entire model. This typically includes:

• The model's architecture (layers, connections).
• The model's learned weights.
• The model's training configuration (specified during compile(), like the optimizer, loss function, and metrics).
• The state of the optimizer (allowing you to resume training exactly where you left off).

Keras provides a simple save() method for this purpose. The recommended format, especially with Keras 3, is the .keras zip archive format. It's a more modern, efficient, and unified format compared to older options like HDF5 (.h5) or TensorFlow's SavedModel (.tf).

import keras
from keras import layers
import numpy as np

# Define a simple sequential model (example)
model = keras.Sequential(
    [
        keras.Input(shape=(784,)),
        layers.Dense(64, activation="relu"),
        layers.Dense(10, activation="softmax"),
    ]
)

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

# Assume the model has been trained...
# model.fit(x_train, y_train, epochs=5, batch_size=32)

# Save the entire model to a single file in the recommended format
model.save("my_model.keras")

print("Model saved successfully to my_model.keras")

This single command packages everything needed into the my_model.keras file.

Loading the Entire Model

To load a model that was saved using model.save(), you use the keras.models.load_model() function. This function reconstructs the model architecture, loads the weights, and restores the training configuration and optimizer state.

import keras
import numpy as np

# Load the model from the file
loaded_model = keras.models.load_model("my_model.keras")

# Verify the model summary
loaded_model.summary()

# You can now use the loaded model for predictions or continue training
# Example prediction (requires appropriate input data)
# dummy_input = np.random.rand(1, 784)
# predictions = loaded_model.predict(dummy_input)
# print("Prediction from loaded model:", predictions)

# You could even resume training if needed
# loaded_model.fit(x_train_more, y_train_more, epochs=2)

Note that if your model includes custom layers, custom activation functions, or other custom objects not part of the core Keras API, you might need to provide these definitions when loading using the custom_objects argument in load_model(). However, for standard layers and functions, this is handled automatically.

Saving Only the Weights

Sometimes, you only need to save the learned parameters (weights) of the model, not the entire architecture or training configuration. This is useful if:

• You plan to use the same model architecture but want to experiment with different weight initializations or training runs.
• You are performing transfer learning, where you load weights into a potentially modified architecture.
• You want a smaller file containing just the parameters.

You can save only the weights using the save_weights() method. Keras typically saves weights using the HDF5 format, often with a .weights.h5 extension, although other backend-specific formats might be used.

import keras
from keras import layers
import numpy as np

# Define the same simple sequential model (example)
model = keras.Sequential(
    [
        keras.Input(shape=(784,)),
        layers.Dense(64, activation="relu"),
        layers.Dense(10, activation="softmax"),
    ]
)

# Compile the model (needed for training, but not saved with weights)
model.compile(optimizer='rmsprop',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# Assume the model has been trained...
# model.fit(x_train, y_train, epochs=5, batch_size=32)

# Save only the weights
model.save_weights("my_model_weights.weights.h5")

print("Model weights saved successfully to my_model_weights.weights.h5")

Loading Only the Weights

To load weights saved with save_weights(), you first need to create an instance of your model with the exact same architecture as the one whose weights you saved. Then, you use the load_weights() method on this newly instantiated model.

import keras
from keras import layers
import numpy as np

# 1. Recreate the exact same model architecture
new_model = keras.Sequential(
    [
        keras.Input(shape=(784,)),
        layers.Dense(64, activation="relu"),
        layers.Dense(10, activation="softmax"),
    ]
)

# 2. Load the saved weights into the new model instance
new_model.load_weights("my_model_weights.weights.h5")

print("Model weights loaded successfully.")

# Now `new_model` has the learned weights.
# Note: You still need to compile this model if you want to train it further
# or evaluate it using compiled metrics.
new_model.compile(optimizer='rmsprop',
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

# You can now use `new_model` for predictions.
# dummy_input = np.random.rand(1, 784)
# predictions = new_model.predict(dummy_input)
# print("Prediction from model with loaded weights:", predictions)

If you try to load weights into a model with a different architecture (e.g., different layers, different number of units), Keras will usually raise an error.

Choosing the Right Method

• Use model.save() and keras.models.load_model() when you want a self-contained representation of your model, including its architecture, weights, and training setup. This is the most common and recommended approach for deployment, sharing, or resuming training. Use the .keras format.
• Use model.save_weights() and model.load_weights() when you only care about the learned parameters and plan to instantiate the model architecture separately using your Python code. This is common in research settings for loading pre-trained weights into standard or custom architectures (e.g., transfer learning).

Being able to save and load models effectively is an essential skill for managing deep learning projects. It integrates tightly with techniques like ModelCheckpoint for training and forms the bridge between model development and model deployment.