Saving weights only in TensorFlow

Introduction

Saving weights only lets you keep just the learned parts of a model. This saves space and lets you reuse the model structure later.

You want to save disk space by not saving the full model.

You plan to rebuild the model architecture in code and load saved weights.

You want to share only the learned parameters without the model code.

You want to continue training later from saved weights.

Syntax

TensorFlow

model.save_weights('path_to_file')

model.load_weights('path_to_file')

The file path can be a filename or a folder path depending on format.

You must have the model architecture defined before loading weights.

Examples

Save and load weights using HDF5 format file.

TensorFlow

model.save_weights('weights.h5')
model.load_weights('weights.h5')

Save and load weights using TensorFlow checkpoint format.

TensorFlow

model.save_weights('checkpoints/weights')
model.load_weights('checkpoints/weights')

Sample Model

This code trains a simple model, saves only its weights, then creates a new model with the same structure and loads the saved weights. Finally, it evaluates the new model to show the weights were loaded correctly.

TensorFlow

import tensorflow as tf

# Create a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(5, activation='relu', input_shape=(3,)),
    tf.keras.layers.Dense(2, activation='softmax')
])

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

# Dummy data
import numpy as np
x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
y = np.array([0, 1, 1], dtype=np.int32)

# Train the model
model.fit(x, y, epochs=2, verbose=0)

# Save weights only
model.save_weights('my_weights')

# Create a new model with the same architecture
new_model = tf.keras.Sequential([
    tf.keras.layers.Dense(5, activation='relu', input_shape=(3,)),
    tf.keras.layers.Dense(2, activation='softmax')
])

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

# Load saved weights
new_model.load_weights('my_weights')

# Evaluate new model on same data
loss, accuracy = new_model.evaluate(x, y, verbose=0)
print(f'Loss: {loss:.4f}, Accuracy: {accuracy:.4f}')

Important Notes

You must define the model architecture exactly the same before loading weights.

Saving weights only is smaller than saving the full model but needs the model code to rebuild.

Use .h5 extension for HDF5 format or no extension for TensorFlow checkpoint format.

Summary

Saving weights only stores just the learned parameters, not the model structure.

You need to recreate the model architecture before loading weights.

This method is useful for saving space and sharing learned parameters.