Introduction

After training a TensorFlow model, the next step is usually saving it, loading it later, and running inference without retraining. The practical details depend on whether you are using modern Keras saving APIs or older graph/session-based TensorFlow code. In current TensorFlow, the safest path is to save a complete model artifact and load it with the same framework family for inference.

Save a Trained Keras Model Correctly

For modern TensorFlow, save the whole model instead of only weights unless you have a specific reason not to. A full save preserves architecture, weights, and compile-related metadata.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(32,)),
5    tf.keras.layers.Dense(16, activation="relu"),
6    tf.keras.layers.Dense(3, activation="softmax"),
7])
8
9model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
10
11# Example only: normally you would train before saving.
12model.save("saved_model.keras")

The .keras format is the clean modern default for Keras models. You can also save in SavedModel format for broader serving workflows.

Load the Model for Inference

Once saved, load it with tf.keras.models.load_model.

1import tensorflow as tf
2import numpy as np
3
4loaded = tf.keras.models.load_model("saved_model.keras")
5
6x = np.random.rand(1, 32).astype("float32")
7pred = loaded(x, training=False)
8
9print(pred.shape)
10print(pred.numpy())

Using training=False ensures inference behavior, which matters for layers such as dropout and batch normalization.

SavedModel Workflow for Serving

If you plan to deploy to TensorFlow Serving or other graph-based runtimes, save as SavedModel directory.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(1)
6])
7
8model.save("saved_model_dir", save_format="tf")
9loaded = tf.keras.models.load_model("saved_model_dir")

This creates a directory artifact rather than a single file. It is often the right choice for production serving systems.

If You Only Saved Weights

If you used model.save_weights(...), you must rebuild the architecture in code before loading.

1import tensorflow as tf
2
3def build_model():
4    model = tf.keras.Sequential([
5        tf.keras.layers.Input(shape=(32,)),
6        tf.keras.layers.Dense(16, activation="relu"),
7        tf.keras.layers.Dense(3, activation="softmax"),
8    ])
9    return model
10
11model = build_model()
12model.load_weights("weights_only.weights.h5")

This is useful in some experiment pipelines, but it is more fragile than saving the complete model.

Prediction Pipeline Hygiene

Loading the model is only part of inference. You also need the same preprocessing logic used during training. If your training pipeline normalized, resized, tokenized, or reordered features, the inference path must do the same.

A clean approach is to keep preprocessing in one function or make it part of the exported model pipeline when possible.

1import numpy as np
2
3def preprocess(x: np.ndarray) -> np.ndarray:
4    return x.astype("float32") / 255.0

Inconsistent preprocessing is one of the most common reasons a “working” loaded model gives bad predictions.

Old Graph-Based TensorFlow Code

If you are dealing with older TensorFlow 1.x graph workflows, loading often involves sessions and graph restoration. That codebase style is still maintained in some legacy systems, but for new work you should prefer TensorFlow 2 plus Keras APIs.

If a project still uses frozen graphs or checkpoint-plus-meta-graph files, it is worth planning a migration rather than expanding the legacy pattern further.

Verification After Load

After reloading a model, verify more than just file existence:

• output shape matches expectation.
• predictions are numerically close to pre-save predictions.
• preprocessing path is identical.
• required custom objects are available at load time if used.

A short regression check prevents subtle serving bugs from slipping into deployment.

Common Pitfalls

• Saving only weights, then forgetting to reconstruct the model architecture correctly.
• Loading a model and skipping training-time preprocessing steps.
• Using inference with training=True by accident.
• Mixing TensorFlow and Keras versions that serialize objects differently.
• Treating legacy graph/session code as the default for new projects.

Summary

• Save the full model when possible, not just weights.
• Load with tf.keras.models.load_model for modern TensorFlow workflows.
• Use SavedModel when serving infrastructure expects directory-based exports.
• Keep inference preprocessing identical to training preprocessing.
• Validate predictions after load instead of assuming serialization worked.