Save / Load

import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers

# 1) Build a model
model = keras.Sequential([
    layers.Input(shape=(28, 28)),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.2),
    layers.Dense(10),
])
model.compile(
    optimizer='adam',
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['sparse_categorical_accuracy'],
)
model.fit(x_train, y_train, epochs=5, validation_split=0.1)

# 2) Save — the modern way (single .keras file; Keras 3 default)
model.save('model.keras')                              # entire model: architecture + weights + optimiser state
loaded = keras.models.load_model('model.keras')
loaded.evaluate(x_test, y_test)

# 3) SavedModel — TF's directory format (best for serving)
model.export('saved_model')                            # creates a directory with assets/, variables/, saved_model.pb
loaded = tf.saved_model.load('saved_model')
loaded.signatures['serving_default']                    # the serving function

# Use SavedModel for:
#   • TensorFlow Serving deployment
#   • TFLite conversion
//   • TFJS conversion
//   • Different runtime (mobile, browser)

# 4) Legacy HDF5 — single .h5 file (Keras 2 default; still works)
model.save('model.h5')
keras.models.load_model('model.h5')
# Not recommended for new code — .keras supports more model types.

# 5) Weights only — for transfer learning / fine-tuning
model.save_weights('weights.weights.h5')               # extension matters in Keras 3
# or:
model.save_weights('weights.weights.h5')
# Load into ARCHITECTURE you've already built
blank = build_same_architecture()
blank.load_weights('weights.weights.h5')

# 6) Checkpoints during training
ckpt = keras.callbacks.ModelCheckpoint(
    filepath='ckpt/best.keras',
    monitor='val_loss',
    mode='min',
    save_best_only=True,
    save_weights_only=False,                            # save whole model
    verbose=1,
)
model.fit(x_train, y_train, validation_split=0.1, epochs=20, callbacks=[ckpt])

# Use save_freq='epoch' (default) or an integer batch count.

# 7) Training resumption pattern
initial_epoch = 0
try:
    model = keras.models.load_model('ckpt/last.keras')
    initial_epoch = int(model.optimizer.iterations.numpy() // steps_per_epoch)
    print(f'resuming from epoch {initial_epoch}')
except Exception:
    model = build_model()

model.fit(
    x_train, y_train,
    epochs=20,
    initial_epoch=initial_epoch,
    callbacks=[keras.callbacks.ModelCheckpoint('ckpt/last.keras', save_best_only=False)],
)

# 8) Custom layers / losses — register them for save/load
@keras.utils.register_keras_serializable()
class MyActivation(layers.Layer):
    def __init__(self, scale=1.0, **kwargs):
        super().__init__(**kwargs)
        self.scale = scale
    def call(self, inputs):
        return tf.nn.swish(inputs) * self.scale
    def get_config(self):
        return {**super().get_config(), 'scale': self.scale}

# With register_keras_serializable, load_model finds your custom class automatically.
# Otherwise pass custom_objects={'MyActivation': MyActivation} to load_model.

# 9) TFLite — mobile / edge / microcontroller deployment
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model')
converter.optimizations = [tf.lite.Optimize.DEFAULT]   # quantisation
tflite_model = converter.convert()
open('model.tflite', 'wb').write(tflite_model)

# Quantised: 4x smaller, slightly less accurate
# To target a specific dataset's distribution, supply a representative_dataset for int8 quant.

# 10) TFJS — browser inference
# pip install tensorflowjs
# tensorflowjs_converter --input_format=tf_saved_model saved_model/ web_model/
# Load in browser: tf.loadGraphModel('web_model/model.json')

# 11) Saving + loading with custom optimiser state
# Saving an entire model (.keras / SavedModel) includes optimiser state by default.
# Saving weights only DOES NOT — recompile + load weights = fresh optimiser.

# 12) Model versioning
import datetime
version = datetime.datetime.utcnow().strftime('%Y%m%d-%H%M%S')
model.save(f'models/{version}/model.keras')
# or use MLflow / Vertex AI Model Registry for managed versioning.

# 13) TensorFlow Serving signature
# When exporting for serving, specify input signature for clear, documented IO:
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 28, 28], dtype=tf.float32)])
def predict_fn(x):
    return { 'probabilities': tf.nn.softmax(model(x)) }

tf.saved_model.save(
    model,
    'saved_model_with_sig',
    signatures={'serving_default': predict_fn},
)

# 14) tf.train.Checkpoint — low-level Keras-independent
ckpt = tf.train.Checkpoint(model=model, optimizer=opt, step=tf.Variable(0))
manager = tf.train.CheckpointManager(ckpt, 'ckpts', max_to_keep=3)
manager.save()                                          # save
ckpt.restore(manager.latest_checkpoint).assert_existing_objects_matched()

# Useful for custom training loops + fine-grained control.

# 15) Migration: TF1 SavedModel → modern Keras
# Use tf.compat.v1 loaders + saved_model_cli if you need to inspect legacy models.
# Re-train where possible; TF1 graphs lose information modern Keras expects.

# 16) Common bugs
# • Saving with .h5 then trying to load a custom layer → 'Unknown layer'; use register_keras_serializable
# • SavedModel signatures missing → 'Could not find a default signature'; provide explicit signatures
# • Loading old .h5 with newer Keras → '__layer_call__ deprecation'; convert to .keras
# • Loading large models from network → slow; cache locally with tf.keras.utils.get_file
# • Forgetting model.eval()-like switch → BatchNorm in training mode at inference; Keras handles via 'training' arg
# • Restoring optimiser then changing learning rate — restore won't override scheduler
# • Loading model exported with different TF version → may need migration script
# • TFLite quantization changing accuracy unexpectedly → test on holdout before shipping
# • Saving subclassed models with custom call() but no get_config() → load fails; implement get_config + from_config

model.save('mymodel.keras')           # Keras v3 native
import tensorflow as tf
restored = tf.keras.models.load_model('mymodel.keras')