machine learning
model persistence
model saving
training models
data science

How to save/restore a model after training?

Master System Design with Codemia

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

Introduction

Saving a trained model matters because training is usually the expensive part of the workflow. A good save-and-restore strategy lets you deploy the model, resume training, and reproduce results without retraining from scratch.

Decide What You Need to Save

There are two common persistence goals:

  • save everything needed for immediate reuse
  • save only weights or parameters and rebuild the model in code

The first option is simpler to load later. The second option gives you more control and is common when the model architecture is already part of the source code.

Saving and Restoring in Keras

Keras can save the full model in one file, including architecture and training configuration.

python
1import numpy as np
2from tensorflow import keras
3
4x = np.random.rand(100, 4).astype("float32")
5y = (x.sum(axis=1) > 2).astype("float32")
6
7model = keras.Sequential([
8    keras.layers.Input(shape=(4,)),
9    keras.layers.Dense(8, activation="relu"),
10    keras.layers.Dense(1, activation="sigmoid"),
11])
12
13model.compile(optimizer="adam", loss="binary_crossentropy")
14model.fit(x, y, epochs=2, verbose=0)
15
16model.save("classifier.keras")

Restoring it is straightforward:

python
restored = keras.models.load_model("classifier.keras")
predictions = restored.predict(x[:2], verbose=0)
print(predictions)

This is the most convenient option when you want to use the model again with minimal ceremony.

If you only save weights, you must rebuild the same architecture first:

python
model.save_weights("weights.weights.h5")

Later:

python
1restored = keras.Sequential([
2    keras.layers.Input(shape=(4,)),
3    keras.layers.Dense(8, activation="relu"),
4    keras.layers.Dense(1, activation="sigmoid"),
5])
6
7restored.compile(optimizer="adam", loss="binary_crossentropy")
8restored.load_weights("weights.weights.h5")

That works, but only if the restored model matches the original architecture.

Saving and Restoring in PyTorch

PyTorch usually saves the model state dictionary rather than the whole Python object. That keeps the persistence format explicit.

python
1import torch
2import torch.nn as nn
3
4
5class Net(nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.layers = nn.Sequential(
9            nn.Linear(4, 8),
10            nn.ReLU(),
11            nn.Linear(8, 1),
12        )
13
14    def forward(self, x):
15        return self.layers(x)
16
17
18model = Net()
19torch.save(model.state_dict(), "model.pt")

To restore:

python
restored = Net()
restored.load_state_dict(torch.load("model.pt", weights_only=True))
restored.eval()

The call to eval() matters for inference because it switches layers such as dropout and batch normalization into evaluation mode.

Save More Than the Bare Model When Needed

Sometimes the model parameters are not enough. You may also need:

  • label encoders
  • tokenizers or vocabularies
  • feature scalers
  • preprocessing pipelines
  • optimizer state if training will resume

For example, a classifier without the same preprocessing logic is not really the same model in production. The model file and the surrounding metadata should travel together.

Resuming Training Is Different from Inference

If you want to continue training later, include whatever state your framework needs for that job. In Keras, a full model save often captures more context than weights alone. In PyTorch, you often save a training checkpoint dictionary:

python
1torch.save({
2    "model_state": model.state_dict(),
3    "optimizer_state": optimizer.state_dict(),
4    "epoch": 5,
5}, "checkpoint.pt")

That is different from a pure inference artifact. It is designed for recovery and continuation.

Common Pitfalls

  • Saving only weights and then forgetting to rebuild the exact same architecture before loading them.
  • Restoring a model but not restoring preprocessing artifacts such as tokenizers or scalers.
  • Using a training checkpoint as if it were a clean deployment artifact without understanding what is inside it.
  • Forgetting model.eval() in PyTorch before inference.
  • Changing library versions or custom layer code and expecting old saved models to load without compatibility checks.

Summary

  • Save models so you can deploy, resume training, and reproduce results without retraining.
  • In Keras, model.save is the simplest full-model workflow.
  • In PyTorch, saving and loading state_dict is the standard explicit approach.
  • For resumed training, save optimizer state and checkpoint metadata, not just model parameters.
  • A usable restored model often requires preprocessing artifacts alongside the model file itself.

Course illustration
Course illustration

All Rights Reserved.