Introduction

Saving a TensorFlow model means preserving enough information to restore it later for inference, continued training, or export. In everyday TensorFlow 2 code, the most common choices are saving the full Keras model or saving only the weights. The right format depends on whether you need the complete model structure or just the learned parameters.

Save the Whole Model

For tf.keras models, the simplest full-model approach is model.save(...).

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(8, activation="relu"),
6    tf.keras.layers.Dense(1)
7])
8
9model.compile(optimizer="adam", loss="mse")
10model.fit(tf.random.normal((10, 4)), tf.random.normal((10, 1)), epochs=1, verbose=0)
11
12model.save("my_model.keras")

This writes a serialized model file that can later be loaded back with:

loaded = tf.keras.models.load_model("my_model.keras")
print(loaded.predict(tf.random.normal((1, 4)), verbose=0))

Saving the full model is convenient because it preserves architecture, weights, and enough metadata to restore the model object directly.

Save Only the Weights

Sometimes you want to recreate the model in code and only save the learned parameters. In that case, use save_weights:

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(8, activation="relu"),
6    tf.keras.layers.Dense(1)
7])
8
9model.save_weights("weights.weights.h5")

To load them later, you must rebuild the same model structure first:

1recreated = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(4,)),
3    tf.keras.layers.Dense(8, activation="relu"),
4    tf.keras.layers.Dense(1)
5])
6
7recreated.load_weights("weights.weights.h5")

This is a good choice when the model definition is version-controlled in code and you only need parameter persistence.

Full Model Versus Weights Only

A practical decision rule is:

• save the full model when you want easy restoration and deployment
• save only weights when model structure is already fixed in code and you want a lighter artifact

The full model is usually easier for inference workflows. Weight-only saving is common in training experiments and checkpoints.

Use Checkpoints During Training

If you need periodic saving during training, use a callback:

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(8, activation="relu"),
6    tf.keras.layers.Dense(1)
7])
8
9model.compile(optimizer="adam", loss="mse")
10
11checkpoint = tf.keras.callbacks.ModelCheckpoint(
12    filepath="checkpoint.weights.h5",
13    save_weights_only=True,
14    save_best_only=True,
15    monitor="loss"
16)
17
18x = tf.random.normal((32, 4))
19y = tf.random.normal((32, 1))
20model.fit(x, y, epochs=3, callbacks=[checkpoint], verbose=0)

This is better than manually saving after every epoch in custom code.

Think About Portability

If the model uses custom layers, losses, or metrics, loading the full model may require those custom objects to be available at restore time. That is one reason some teams prefer weight-only checkpoints plus explicit model-construction code.

Also, if your real goal is deployment to TensorFlow Serving or another serving tool, you may want an export-oriented format rather than a training-oriented checkpoint routine. The storage choice should match the downstream use case.

Separate Training Checkpoints from Deployment Exports

A useful habit is to distinguish between files meant for continuing training and files meant for serving or shipping a model. Checkpoints are great during experimentation because they are lightweight and frequent. Full model saves are usually better when you want a self-contained artifact that another process can load later without rebuilding the architecture manually.

That distinction keeps project directories cleaner and makes it easier to reason about which files are safe to overwrite and which ones represent a publishable model version.

Common Pitfalls

• Saving only weights and expecting load_model to reconstruct the full model.
• Loading weights into a model with a different architecture.
• Forgetting that custom layers or custom losses may affect how full-model loading works.
• Mixing training checkpoints and deployment artifacts without a clear separation.
• Assuming model saving is version-agnostic across arbitrary TensorFlow environments.

Summary

• Use model.save(...) when you want to persist the whole Keras model.
• Use save_weights(...) when you only need the learned parameters.
• Restore full models with load_model(...) and weights with load_weights(...).
• Use checkpoint callbacks for repeated saves during training.
• Choose the save format based on whether the goal is training continuation, inference, or deployment.