Keras 3 .keras save format: preserving optimizer state during zero‑downtime migration

Answer

Yes – model.save('model.keras') in Keras 3 stores the full optimizer state (config and variable values) but it does not embed the training history (epoch‑level loss/metrics). Custom layers, metrics, or losses that were defined in Keras 2 must be re‑registered with keras.saving.register_keras_serializable or supplied via the custom_objects argument when loading the .keras file in the same TensorFlow/Keras version.

Confirmed Facts (from the Keras 3 docs and release notes)

• Architectural data, trainable weights, and the optimizer’s variables are written into the ZIP archive under optimizer/variables and optimizer/config.
• The history object is not part of the file; it must be saved separately (e.g., via pickle or a JSON export).
• When the same backend (TensorFlow) and compatible Keras version are used to load the model, the optimizer can be resumed immediately and fine‑tuning can continue from the exact point of last training.
• Custom objects that were present in a Keras 2 checkpoint need to be registered again in Keras 3; otherwise deserialization will raise a ValueError.

Steps for Zero‑Downtime Migration

1. Save the current checkpoint as a .keras file (do this in a separate process that does not block the inference service):

   # In a temporary script or interactive shell
   import tensorflow as tf
   from tensorflow import keras
   
   model = keras.models.load_model('model.h5')  # Keras 2 checkpoint
   model.save('model.keras')  # Keras 3 format
   

2. Verify optimizer weights are preserved by loading both models in the same environment and comparing the weights:

   # In the same Python process
   from tensorflow import keras
   
   orig = keras.models.load_model('model.h5')
   new  = keras.models.load_model('model.keras')
   
   orig_opt = orig.optimizer.get_weights()
   new_opt  = new.optimizer.get_weights()
   
   # Simple element‑wise check
   import numpy as np
   print(np.allclose(orig_opt, new_opt, atol=1e-6))
   

3. Register any custom objects that were used in the original model. Example for a custom layer:

   from tensorflow import keras
   
   @keras.saving.register_keras_serializable(package='Custom', name='MyLayer')
   class MyLayer(keras.layers.Layer):
       ...
   

   Alternatively, pass them as custom_objects to load_model.
4. Deploy the new .keras file to the inference service. Because the optimizer state is present, you can immediately resume fine‑tuning if needed.
5. Optional: Persist training history if you need to resume epoch counts:

   # Save history separately before migration
   import json
   with open('history.json', 'w') as f:
       json.dump(history.history, f)
   

How to Verify Optimizer State Without Retraining

1. Load the original and the new model in the same environment.
2. Compare the optimizer variable tensors (shown above with np.allclose).
3. Optionally, run model.evaluate(test_data) on both models and confirm that the loss and metrics match the last checkpoint values within a small tolerance.

Diagnostic Question

Did the original training use a custom training loop (e.g., tf.GradientTape) instead of model.fit? If so, the optimizer state may not be fully captured by .keras, and you might need to manually export and import the optimizer weights.

Summary

Converting a Keras 2 HDF5 checkpoint to the Keras 3 .keras ZIP format preserves the optimizer variables, enabling immediate resumption of fine‑tuning. The training history must be saved separately. Custom objects require re‑registration for successful deserialization.