Saving weights only in TensorFlow - ML Experiment: Train & Evaluate

Experiment - Saving weights only

Problem:You have trained a neural network model on a dataset, but the model file size is large because it saves the entire model architecture and optimizer state. You want to save only the model weights to reduce storage size and load them later into the same model architecture.

Current Metrics:Training accuracy: 92%, Validation accuracy: 89%, Model file size: 25MB

Issue:Saving the full model results in a large file size, which is inefficient for storage and sharing. You want to save only the weights without losing the ability to restore the model later.

Your Task

Save only the model weights after training and then load these weights into the same model architecture to make predictions. Confirm that the loaded model produces the same validation accuracy.

You must keep the model architecture code unchanged.

You cannot save the entire model (architecture + weights + optimizer).

You must use TensorFlow's recommended methods for saving and loading weights.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

TensorFlow

import tensorflow as tf
from tensorflow.keras import layers, models

# Define the model architecture
model = models.Sequential([
    layers.Dense(64, activation='relu', input_shape=(20,)),
    layers.Dense(10, activation='softmax')
])

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

# Generate dummy data for training and validation
import numpy as np
X_train = np.random.random((1000, 20))
y_train = np.random.randint(0, 10, 1000)
X_val = np.random.random((200, 20))
y_val = np.random.randint(0, 10, 200)

# Train the model
model.fit(X_train, y_train, epochs=5, batch_size=32, validation_data=(X_val, y_val))

# Save only the weights
weights_file = 'model_weights.weights.h5'
model.save_weights(weights_file)

# Create a new model instance with the same architecture
new_model = models.Sequential([
    layers.Dense(64, activation='relu', input_shape=(20,)),
    layers.Dense(10, activation='softmax')
])

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

# Load the saved weights into the new model
new_model.load_weights(weights_file)

# Evaluate the new model on validation data
loss, accuracy = new_model.evaluate(X_val, y_val, verbose=0)
print(f'Validation accuracy after loading weights: {accuracy * 100:.2f}%')

Used model.save_weights() to save only the weights to a file.

Created a new model instance with the same architecture.

Used model.load_weights() to load the saved weights into the new model.

Evaluated the new model to confirm it matches the original validation accuracy.

Changed weights file name to end with '.weights.h5' to comply with TensorFlow's recommended naming convention for saving weights in HDF5 format.

Results Interpretation

Before saving weights only: Model file size was 25MB including architecture and optimizer state.

After saving weights only: Weights file size is significantly smaller (usually a few MBs), and validation accuracy remains at 89% after loading weights into the same architecture.

Saving only the model weights reduces storage size and allows you to restore the model later by defining the same architecture and loading the weights. This is useful when you want to share or store models efficiently without saving the entire model configuration.

Bonus Experiment

Try saving and loading weights using the TensorFlow checkpoint format instead of HDF5 format.

💡 Hint

Use model.save_weights('path', save_format='tf') to save in TensorFlow format and load_weights() similarly.