How to save and load Keras model weights

A trained network often needs to hand its learned parameters to a fresh process without carrying the optimizer or model configuration. A weights-only artifact keeps that handoff small while leaving architecture construction under application code.

In Keras 3, model.save_weights() writes a single file whose name ends in .weights.h5 by default. load_weights() restores values according to the network topology, so the receiving model must create compatible weighted layers in the same order and with the same shapes.

A fixed two-feature weight fixture makes the round trip deterministic without depending on training noise. The receiving model comes from the same factory, and a fail-capable NumPy assertion compares its prediction with the source model's prediction after restoration.

1. Create weights_roundtrip.py with the JAX backend selection and reusable model factory.

   weights_roundtrip.py

   import os
    
   os.environ.setdefault("KERAS_BACKEND", "jax")
    
   from pathlib import Path
    
   import keras
   import numpy as np
    
    
   def build_model():
       return keras.Sequential(
           [
               keras.Input(
                   shape=(2,),
                   name="features",
               ),
               keras.layers.Dense(
                   1,
                   name="score",
               ),
           ]
       )

2. Append the source-weight persistence block below build_model().

   weights_roundtrip.py

   weights_path = Path("score.weights.h5")
   sample = np.array(
       [[2.0, 4.0]],
       dtype="float32",
   )
   source_model = build_model()
   source_model.layers[0].set_weights(
       [
           np.array(
               [[0.25], [0.75]],
               dtype="float32",
           ),
           np.array([-0.1], dtype="float32"),
       ]
   )
   prediction_before = np.asarray(
       source_model(sample, training=False)
   )
   source_model.save_weights(weights_path)

   The .weights.h5 suffix selects the single-file weights format. The receiving model requires the same layer order and tensor shapes.

3. Add the restore and prediction comparison below save_weights().

   weights_roundtrip.py

   restored_model = build_model()
   restored_model.load_weights(weights_path)
   prediction_after = np.asarray(
       restored_model(sample, training=False)
   )
   np.testing.assert_allclose(
       prediction_before,
       prediction_after,
       atol=1e-7,
   )
    
   print(f"weights file: {weights_path}")
   print(
       "prediction before save: "
       f"{prediction_before[0, 0]:.6f}"
   )
   print(
       "prediction after load: "
       f"{prediction_after[0, 0]:.6f}"
   )
   print("predictions match: True")

4. Inspect the assembled weights_roundtrip.py file for the three constructed sections.

   weights_roundtrip.py

   import os
    
   os.environ.setdefault("KERAS_BACKEND", "jax")
    
   from pathlib import Path
    
   import keras
   import numpy as np
    
    
   def build_model():
       return keras.Sequential(
           [
               keras.Input(
                   shape=(2,),
                   name="features",
               ),
               keras.layers.Dense(
                   1,
                   name="score",
               ),
           ]
       )
    
    
   weights_path = Path("score.weights.h5")
   sample = np.array(
       [[2.0, 4.0]],
       dtype="float32",
   )
   source_model = build_model()
   source_model.layers[0].set_weights(
       [
           np.array(
               [[0.25], [0.75]],
               dtype="float32",
           ),
           np.array([-0.1], dtype="float32"),
       ]
   )
   prediction_before = np.asarray(
       source_model(sample, training=False)
   )
   source_model.save_weights(weights_path)
    
    
   restored_model = build_model()
   restored_model.load_weights(weights_path)
   prediction_after = np.asarray(
       restored_model(sample, training=False)
   )
   np.testing.assert_allclose(
       prediction_before,
       prediction_after,
       atol=1e-7,
   )
    
   print(f"weights file: {weights_path}")
   print(
       "prediction before save: "
       f"{prediction_before[0, 0]:.6f}"
   )
   print(
       "prediction after load: "
       f"{prediction_after[0, 0]:.6f}"
   )
   print("predictions match: True")

5. Run weights_roundtrip.py to verify the restored prediction matches the source prediction.

   $ python weights_roundtrip.py
   weights file: score.weights.h5
   prediction before save: 3.400000
   prediction after load: 3.400000
   predictions match: True

   The assertion stops execution if the restored prediction differs, while predictions match: True confirms the saved values produced the same output in a separately constructed model.