Save & Restore

Training is expensive, so saving is how you keep the result. Keras gives you one format that stores architecture, weights, and optimizer state together, and one that stores weights alone. Knowing which to use — and what a checkpoint does not include — is the difference between a reproducible model and a broken deploy.

The idea

There are three things you might save, and they are not the same:

• Weights only — the arrays. Small and fast, but meaningless without the code that defines the architecture.
• Whole model — architecture plus weights plus optimizer state plus the compile configuration. Self-contained.
• Checkpoint during training — the same as the whole model, written periodically so a crash does not cost you the run.

The Keras v3 format (.keras) is the current whole-model format: a zip archive holding a JSON config and an HDF5 weights file. The older .h5 whole-model format is legacy and silently drops anything it cannot serialize, which includes subclassed models and custom layers. tf.train.Saver and .ckpt files belong to TF 1.x graph mode.

Worked example

Sizes for the small flatten-and-dense MNIST model, which has 784 × 128 + 128 + 128 × 10 + 10 = 101,770 parameters:

Float32 is 4 bytes per value, so the weights are 101,770 × 4 ≈ 407 KB. Adam roughly doubles that because it keeps a first and second moment for every trainable variable, which is why a full checkpoint is noticeably larger than the weights file.

In code

Train, keep the best checkpoint, then reload for inference:

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28, 1)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation="relu"),
    tf.keras.layers.Dense(10),
])
model.compile(optimizer="adam",
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=["accuracy"])

cb = tf.keras.callbacks.ModelCheckpoint(
    "best.keras", monitor="val_accuracy", save_best_only=True)
model.fit(x_train, y_train, epochs=5, validation_split=0.1, callbacks=[cb])

restored = tf.keras.models.load_model("best.keras")
probs = tf.nn.softmax(restored.predict(x_test))

If you save weights only, rebuild the identical architecture before loading, or the variable order will not match:

model.save_weights("mnist.weights.h5")
model.load_weights("mnist.weights.h5")

load_model fails on a custom layer unless you pass custom_objects={"MyLayer": MyLayer} or decorate the class with @tf.keras.utils.register_keras_serializable.

Check yourself

1. Why can a checkpoint be larger than the weights file for the same model?
2. What must exist in code before load_weights will succeed?
3. When would you use ModelCheckpoint(save_best_only=True) instead of one final save at the end of training?

Key takeaways

• .keras is the current whole-model format; whole-model .h5 saving is legacy.
• Weights-only files are portable but require the architecture to be rebuilt and matched.
• The optimizer state is part of the checkpoint — losing it makes a resumed run restart poorly.