TensorFlow Saving and Loading Models
Training a model takes time, data, and computing power. Once training finishes, you need to save the model so you can reuse it later for predictions, share it with teammates, or deploy it to a production server. TensorFlow provides multiple formats for saving models, each suited to different use cases. This topic explains every format, when to use each, and how to reload a saved model correctly.
What Needs to Be Saved
A complete model save includes: 1. Architecture — the layer structure and connections 2. Weights — all the learned numbers (kernel and bias values) 3. Optimizer state — learning rate and momentum values (for resuming training) 4. Compile config — loss, optimizer name, metrics (for further training)
Format 1 — Keras Native Format (.keras)
The .keras format is the modern recommended format for saving and loading Keras models. It stores everything in a single file: architecture, weights, optimizer state, and compile config.
import tensorflow as tf
# Save
model.save('my_model.keras')
# Load
loaded_model = tf.keras.models.load_model('my_model.keras')
# Verify the loaded model works
results = loaded_model.evaluate(x_test, y_test)
print("Loaded model accuracy:", results[1])
Format 2 — SavedModel Format
SavedModel is TensorFlow's universal format. It saves not just the Keras model but the full TensorFlow computation graph, making it compatible with TensorFlow Serving, TFLite conversion, TensorFlow.js export, and multi-language deployment.
# Save as SavedModel directory
model.save('saved_model_dir/')
# Creates: saved_model_dir/
# ├── saved_model.pb (graph definition)
# └── variables/
# ├── variables.index
# └── variables.data-00000-of-00001
# Load
loaded_model = tf.keras.models.load_model('saved_model_dir/')
Format 3 — HDF5 Format (.h5)
The HDF5 format (legacy) saves architecture and weights in a single .h5 file. It works for most models but does not support custom objects as well as the .keras format. New projects should prefer .keras.
# Save
model.save('my_model.h5')
# Load
loaded_model = tf.keras.models.load_model('my_model.h5')
Saving Only the Weights
Sometimes you want to save just the learned weights without the architecture, for example when you plan to rebuild the same model in code and just restore the parameters.
# Save weights only
model.save_weights('weights_only.weights.h5')
# Load weights into a model with the same architecture
new_model = build_model() # Create the same architecture
new_model.load_weights('weights_only.weights.h5')
Checkpointing During Training
A checkpoint saves the model's weights at specific points during training. This protects against crashes and lets you restore the best-performing model rather than the final one.
import tensorflow as tf
# Save checkpoint at the end of every epoch
checkpoint_cb = tf.keras.callbacks.ModelCheckpoint(
filepath='checkpoints/model_epoch_{epoch:02d}.keras',
save_weights_only=False,
save_freq='epoch',
verbose=1
)
model.fit(x_train, y_train, epochs=20, callbacks=[checkpoint_cb])
Save Only the Best Model
best_model_cb = tf.keras.callbacks.ModelCheckpoint(
filepath='best_model.keras',
monitor='val_accuracy', # Watch this metric
save_best_only=True, # Only overwrite if this metric improves
mode='max', # Higher val_accuracy is better
verbose=1
)
model.fit(x_train, y_train,
epochs=50,
validation_data=(x_val, y_val),
callbacks=[best_model_cb])
# After training, load the best model (not necessarily the last epoch)
best_model = tf.keras.models.load_model('best_model.keras')
TensorFlow Checkpoints (Low-Level)
For custom training loops that do not use model.fit(), use TensorFlow's low-level checkpoint system:
import tensorflow as tf
optimizer = tf.keras.optimizers.Adam()
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
manager = tf.train.CheckpointManager(checkpoint,
directory='./tf_checkpoints',
max_to_keep=3)
# Save
save_path = manager.save()
print(f"Saved checkpoint: {save_path}")
# Restore latest checkpoint
checkpoint.restore(manager.latest_checkpoint)
print("Restored from:", manager.latest_checkpoint)
Saving and Loading Custom Models
If your model uses custom layers, custom loss functions, or custom metrics, you must pass a custom_objects dictionary when loading:
# Custom activation function
def my_activation(x):
return x * tf.nn.sigmoid(x) # Swish
# Save works normally
model.save('custom_model.keras')
# Load requires custom_objects
loaded = tf.keras.models.load_model(
'custom_model.keras',
custom_objects={'my_activation': my_activation}
)
Format Comparison
Format Extension Architecture Weights Optimizer Deployment ────────────────────────────────────────────────────────────────────────── Keras native .keras Yes Yes Yes Keras only SavedModel directory Yes Yes Yes Universal HDF5 (legacy) .h5 Yes Yes Yes Keras only Weights only .weights.h5 No Yes No Keras only TF Checkpoint directory No Yes Yes TF only ────────────────────────────────────────────────────────────────────────── Use .keras for most projects. Use SavedModel when deploying to TF Serving or converting to TFLite. Use Weights only when sharing parameters across experiments.
Converting SavedModel for Deployment
# Convert SavedModel to TFLite for mobile
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model_dir/')
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
# Convert SavedModel to TensorFlow.js for browser
# (requires tensorflowjs package)
# tensorflowjs_converter --input_format=tf_saved_model saved_model_dir/ tfjs_dir/
Loading for Inference Only (No Training)
# Load model and run predictions — no need to recompile
loaded_model = tf.keras.models.load_model('my_model.keras')
# Single prediction
sample = x_test[0:1] # Keep batch dimension: shape (1, features)
prediction = loaded_model.predict(sample)
print("Prediction:", prediction)
# Batch prediction
all_predictions = loaded_model.predict(x_test, batch_size=64)
Verifying a Loaded Model
# Always verify the loaded model matches the original
original_results = model.evaluate(x_test, y_test, verbose=0)
loaded_results = loaded_model.evaluate(x_test, y_test, verbose=0)
print(f"Original accuracy: {original_results[1]:.4f}")
print(f"Loaded accuracy: {loaded_results[1]:.4f}")
# These should be identical
Saving and loading models correctly is an essential production skill. A missing custom_objects argument, a wrong file format for deployment, or failing to save the best checkpoint rather than the last one are common mistakes that waste hours of training time. The next topic moves into data handling — how to build efficient data pipelines that feed models without bottlenecking on slow disk reads.