TensorFlow
Keras
Model Saving
Model Restoration
Deep Learning

How to save/restore large model in tensorflow 2.0 w/ keras?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Saving a large Keras model is not just a matter of writing one file to disk. You usually need two different behaviors: reliable periodic checkpoints during training and a clean artifact for later loading or deployment. In modern TensorFlow and Keras, the most practical strategy is to checkpoint frequently, save a final full model explicitly, and keep preprocessing assets alongside the model.

Pick the Right Save Strategy

There are three common save targets, and they are not interchangeable:

  • weights only
  • full model
  • periodic checkpoints during training

If you want the easiest reload for inference, save the full model. If you want crash recovery during a long training run, use checkpoints. If you want a lightweight artifact and you already control the architecture code, weights-only saving is fine.

Save a Full Keras Model

For a complete restore path, use the Keras model format:

python
1import tensorflow as tf
2from tensorflow import keras
3
4model = keras.Sequential([
5    keras.layers.Input(shape=(128,)),
6    keras.layers.Dense(256, activation="relu"),
7    keras.layers.Dense(10, activation="softmax"),
8])
9
10model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
11model.save("artifacts/classifier.keras")

Load it later with:

python
restored = keras.models.load_model("artifacts/classifier.keras")

This is the simplest option when you want to preserve architecture, weights, and training configuration together.

Use Checkpoints During Long Training Runs

Large models often train for hours or days, so waiting until the end is risky. ModelCheckpoint lets you save progress periodically.

python
1checkpoint_cb = keras.callbacks.ModelCheckpoint(
2    filepath="checkpoints/epoch_{epoch:02d}.weights.h5",
3    save_weights_only=True,
4    save_freq="epoch"
5)
6
7model.fit(
8    x_train,
9    y_train,
10    epochs=20,
11    callbacks=[checkpoint_cb]
12)

That produces a sequence of weight files you can use for recovery if the job is interrupted. For very large models, weights-only checkpoints are often more manageable than repeatedly writing a full model artifact.

Restore the Latest Checkpoint

To resume training, rebuild the model structure, compile it, and load the most recent checkpoint:

python
1import os
2import tensorflow as tf
3from tensorflow import keras
4
5model = keras.Sequential([
6    keras.layers.Input(shape=(128,)),
7    keras.layers.Dense(256, activation="relu"),
8    keras.layers.Dense(10, activation="softmax"),
9])
10
11model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
12
13latest = tf.train.latest_checkpoint("checkpoints")
14if latest:
15    model.load_weights(latest)
16    print("Loaded", latest)

This is the right recovery flow when you are continuing a large training job rather than simply doing inference.

Keep Preprocessing Assets with the Model

A model file alone is often not enough. Many restore failures are really preprocessing failures. Save the tokenizer, vocabulary, feature normalization statistics, label mapping, and any configuration needed to turn raw input into model-ready tensors.

For example:

python
1import json
2
3metadata = {
4    "labels": ["cat", "dog", "horse"],
5    "max_length": 128
6}
7
8with open("artifacts/metadata.json", "w", encoding="utf-8") as f:
9    json.dump(metadata, f)

If those assets are missing, the model may load successfully and still produce meaningless results.

Storage and Retention Matter for Large Models

Large checkpoints can consume disk quickly. A sensible retention policy matters:

  • keep the latest checkpoint
  • keep the best checkpoint by validation metric
  • delete older intermediate checkpoints
  • save one final deployable artifact separately

Do not treat training checkpoints as your only deployment artifact. The file you want for resuming training is not always the one you want to ship to production.

Common Pitfalls

One common mistake is saving weights only and later forgetting the exact architecture needed to reload them. That works only if the model code is recreated exactly.

Another mistake is writing every checkpoint forever. Large models can exhaust disk or object storage much faster than expected.

Developers also often forget the non-model assets required for inference, such as vocabularies, normalization constants, and label mappings.

Finally, avoid assuming that a single save format solves training recovery, inference portability, and deployment equally well. In practice, you usually want checkpoints for training and a separate final artifact for serving.

Summary

  • Use full-model saves for the easiest restore path.
  • Use periodic checkpoints to protect long-running training jobs.
  • Use weights-only saves when architecture is managed separately.
  • Save preprocessing metadata with the model, not just the weights.
  • For large models, retention policy and artifact purpose matter as much as the save API itself.

Course illustration
Course illustration

All Rights Reserved.