TensorFlow Model Saving and Loading

In machine learning and deep learning projects, model saving and loading are crucial steps.

TensorFlow provides multiple ways to save and restore models, enabling developers to:

  • Save trained models for later use
  • Share models with other developers
  • Resume training from checkpoints
  • Deploy models to production environments

TensorFlow 2.x mainly supports three model saving formats:

  1. SavedModel format (recommended)
  2. HDF5 format (.h5)
  3. Legacy Keras format

Saving the Entire Model

SavedModel Format

SavedModel is the model saving format recommended by TensorFlow. It contains complete model information:

Example

import tensorflow as tf

# Create and train a simple model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)

# Save as SavedModel format
model.save('my_model')  # Note: no file extension

Directory structure after saving:

my_model/
├── assets/
├── variables/
│   ├── variables.data-00000-of-00001
│   └── variables.index
└── saved_model.pb

HDF5 Format

HDF5 is another commonly used model saving format:

Example

# Save as HDF5 format
model.save('my_model.h5')  # Note the .h5 extension

Differences Between the Two Formats

Feature SavedModel HDF5
Contains custom objects Yes Requires extra configuration
Contains optimizer state Yes Optional
TensorFlow Serving Native support Not supported
File size Larger Smaller

Loading the Entire Model

Loading from SavedModel

Example

# Load from SavedModel
loaded_model = tf.keras.models.load_model('my_model')

# Verify model
loss, acc = loaded_model.evaluate(x_test, y_test, verbose=2)
print(f"Restored model, accuracy: {100*acc:.1f}%")

Loading from HDF5 File

Example

# Load from HDF5 file
loaded_model = tf.keras.models.load_model('my_model.h5')

# Verify model
loss, acc = loaded_model.evaluate(x_test, y_test, verbose=2)
print(f"Restored model, accuracy: {100*acc:.1f}%")

Selective Saving and Loading

Saving Only Weights

Example

# Save weights
model.save_weights('my_model_weights')

# Save weights in HDF5 format
model.save_weights('my_model_weights.h5')

Loading Weights

Example

# Create a model with the same architecture
new_model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
new_model.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])

# Load weights
new_model.load_weights('my_model_weights')

# Or for .h5 files
new_model.load_weights('my_model_weights.h5')

Saving Checkpoints for Custom Training Loops

Example

# Create checkpoint callback
checkpoint_path = "training_1/cp.ckpt"
cp_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    save_weights_only=True,
    verbose=1)

# Train model using callback
model.fit(x_train, y_train,
          epochs=10,
          callbacks=[cp_callback])

Best Practices for Model Saving and Loading

  1. Production deployment: Prioritize using the SavedModel format
  2. Cross-platform sharing: HDF5 format is more universal
  3. Recovering from training interruption: Use checkpoint callbacks to save periodically
  4. Handling custom objects： 
    model.save('custom_model', save_format='tf')
  5. Model version control: Create different directories for different versions of models

Common Problems and Solutions

Custom Layer/Model Saving Issues

Example

# Custom layer example
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super().__init__(**kwargs)
        self.units = units
   
    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True)
   
    def call(self, inputs):
        return tf.matmul(inputs, self.w)
   
    def get_config(self):
        config = super().get_config()
        config.update({"units": self.units})
        return config

# Use custom layer and save
model = tf.keras.Sequential([CustomLayer(10)])
model.compile(optimizer='adam', loss='mse')
model.save('custom_model')  # Custom layer will be saved automatically

Cross-Version Compatibility Issues

  • Try to use the same version of TensorFlow to save and load models
  • For production environments, consider using TensorFlow Serving to avoid version issues

Optimizing Large Model Saving

Example

# Use save_weights instead of save to reduce saving time
model.save_weights('large_model_weights.h5')
Other Extensions