Practice Now
Keras
LSTM
model saving
model restoration
neural networks

How to save and restore Keras LSTM model?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Deep LearningTensorFlow & Keras

Introduction

Keras provides several ways to save and restore LSTM models: model.save() saves the complete model (architecture, weights, optimizer state) to a single file, model.save_weights() saves only the weights, and model.to_json() exports only the architecture. The recommended modern format is the Keras native format (.keras), which replaced the legacy HDF5 (.h5) format. For TensorFlow deployment, use tf.saved_model.save() to export a SavedModel directory.

Saving the Complete Model

python
1import numpy as np
2from tensorflow import keras
3from tensorflow.keras.models import Sequential
4from tensorflow.keras.layers import LSTM, Dense
5
6# Build an LSTM model
7model = Sequential([
8    LSTM(64, input_shape=(10, 1), return_sequences=True),
9    LSTM(32),
10    Dense(1)
11])
12model.compile(optimizer='adam', loss='mse')
13
14# Train the model
15X_train = np.random.randn(1000, 10, 1)
16y_train = np.random.randn(1000, 1)
17model.fit(X_train, y_train, epochs=5, batch_size=32)
18
19# Save the complete model (Keras 3 native format — recommended)
20model.save('lstm_model.keras')
21
22# Or save in legacy HDF5 format
23model.save('lstm_model.h5')

model.save() stores everything needed to restore the model: layer architecture, trained weights, optimizer state, and compilation config. The .keras format is the default in Keras 3+.

Restoring the Complete Model

python
1from tensorflow.keras.models import load_model
2
3# Load from Keras native format
4restored_model = load_model('lstm_model.keras')
5
6# Or from HDF5
7restored_model = load_model('lstm_model.h5')
8
9# Verify the model works
10predictions = restored_model.predict(X_train[:5])
11print(predictions)
12
13# Continue training
14restored_model.fit(X_train, y_train, epochs=3)

load_model() reconstructs the full model including the optimizer, so you can resume training immediately without recompiling.

Saving and Loading Weights Only

python
1# Save weights only
2model.save_weights('lstm_weights.weights.h5')
3
4# To restore, rebuild the model architecture first
5new_model = Sequential([
6    LSTM(64, input_shape=(10, 1), return_sequences=True),
7    LSTM(32),
8    Dense(1)
9])
10
11# Load weights into the new model
12new_model.load_weights('lstm_weights.weights.h5')
13
14# Must recompile before training (optimizer state is lost)
15new_model.compile(optimizer='adam', loss='mse')

Weight-only saving is useful when the model architecture is defined in code and you only need to store the learned parameters. The architecture must match exactly when loading.

Saving Architecture as JSON

python
1# Save architecture to JSON
2json_config = model.to_json()
3with open('lstm_architecture.json', 'w') as f:
4    f.write(json_config)
5
6# Restore architecture from JSON
7from tensorflow.keras.models import model_from_json
8
9with open('lstm_architecture.json', 'r') as f:
10    json_config = f.read()
11
12restored_architecture = model_from_json(json_config)
13restored_architecture.summary()
14
15# Load weights separately
16restored_architecture.load_weights('lstm_weights.weights.h5')
17restored_architecture.compile(optimizer='adam', loss='mse')

JSON export captures the layer configuration but not weights or optimizer state. This is useful for sharing model designs or version-controlling architecture changes.

SavedModel Format for TensorFlow Serving

python
1import tensorflow as tf
2
3# Export as TensorFlow SavedModel
4tf.saved_model.save(model, 'saved_model_dir/')
5
6# The directory contains:
7# saved_model_dir/
8#   saved_model.pb
9#   variables/
10#     variables.data-00000-of-00001
11#     variables.index
12#   assets/
13
14# Load SavedModel
15loaded = tf.saved_model.load('saved_model_dir/')
16
17# Or load as a Keras model
18loaded_keras = tf.keras.models.load_model('saved_model_dir/')
19loaded_keras.predict(X_train[:5])

SavedModel is the standard format for TensorFlow Serving, TensorFlow Lite conversion, and TensorFlow.js export.

Checkpointing During Training

python
1from tensorflow.keras.callbacks import ModelCheckpoint
2
3# Save the best model during training
4checkpoint = ModelCheckpoint(
5    'best_lstm.keras',
6    monitor='val_loss',
7    save_best_only=True,
8    mode='min',
9    verbose=1
10)
11
12# Save every epoch
13periodic_checkpoint = ModelCheckpoint(
14    'lstm_epoch_{epoch:02d}_loss_{val_loss:.4f}.keras',
15    save_freq='epoch'
16)
17
18model.fit(
19    X_train, y_train,
20    validation_split=0.2,
21    epochs=50,
22    callbacks=[checkpoint, periodic_checkpoint]
23)
24
25# Load the best model after training
26best_model = load_model('best_lstm.keras')

ModelCheckpoint automatically saves the model at specified intervals. save_best_only=True keeps only the version with the lowest validation loss.

Custom LSTM with Custom Objects

python
1from tensorflow.keras.layers import Layer
2
3class AttentionLSTM(Layer):
4    def __init__(self, units, **kwargs):
5        super().__init__(**kwargs)
6        self.units = units
7        self.lstm = LSTM(units, return_sequences=True)
8        self.attention = Dense(1, activation='softmax')
9
10    def call(self, inputs):
11        lstm_out = self.lstm(inputs)
12        weights = self.attention(lstm_out)
13        return tf.reduce_sum(lstm_out * weights, axis=1)
14
15    def get_config(self):
16        config = super().get_config()
17        config.update({'units': self.units})
18        return config
19
20# Save model with custom layer
21model.save('custom_lstm.keras')
22
23# Load with custom_objects
24loaded = load_model('custom_lstm.keras', custom_objects={'AttentionLSTM': AttentionLSTM})

Custom layers must implement get_config() to be serializable. When loading, pass custom_objects so Keras can reconstruct the custom layer.

Common Pitfalls

  • Architecture mismatch when loading weights: load_weights() requires the model architecture to be identical to the one used during saving. Adding, removing, or renaming layers causes shape mismatch errors.
  • Missing custom_objects on load: Models with custom layers, losses, or metrics fail to load without custom_objects. Register them via @keras.utils.register_keras_serializable() or pass them to load_model().
  • HDF5 format deprecated in Keras 3: The .h5 format is legacy. Use .keras for new projects. HDF5 does not support some Keras 3 features like custom saving logic.
  • Optimizer state not saved with save_weights: save_weights() only saves layer weights. Optimizer state (momentum, learning rate schedule) is lost. Use model.save() if you need to resume training exactly where you left off.
  • Large LSTM models and memory: Loading a large LSTM model allocates memory for all layers at once. On memory-constrained systems, use tf.lite.TFLiteConverter to convert to a smaller TensorFlow Lite model before deployment.

Summary

  • Use model.save('model.keras') to save the complete model (architecture + weights + optimizer)
  • Use load_model('model.keras') to restore and resume training
  • Use save_weights() / load_weights() when architecture is defined in code
  • Use ModelCheckpoint callback to save the best model during training
  • Use SavedModel format (tf.saved_model.save) for TensorFlow Serving deployment
  • Custom layers require get_config() implementation and custom_objects on load
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.