TensorFlow Cheatsheet

Saving and Loading

Use this TensorFlow reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

Format Overview

FormatExtensionSavesBest for
Keras native.kerasarchitecture + weights + optimizer + configKeras 3 default; fully reproducible
SavedModeldirectoryarchitecture + weights + computation graphTF Serving, TFLite, TF.js export
HDF5.h5architecture + weights + optimizerLegacy; avoid for new projects
Weights only.weights.h5 (extension required)weights tensors onlyPartial restore, transfer learning

Saving

Keras 3 (TF ≥ 2.16): model.save() requires a .keras or .h5 filename — passing a directory raises. SavedModel is produced with model.export().

# Recommended: Keras native format
model.save('model.keras')

# SavedModel for serving/TFLite — export, not save
model.export('saved_model_dir/')
tf.saved_model.save(model, 'saved_model_dir/')   # low-level equivalent

# HDF5 (legacy)
model.save('model.h5')

# Weights only — filename MUST end in .weights.h5
model.save_weights('weights.weights.h5')

Loading

import tensorflow as tf
from tensorflow import keras

# Keras native
model = keras.models.load_model('model.keras')

# HDF5 (legacy)
model = keras.models.load_model('model.h5')

# SavedModel — load_model CANNOT open a SavedModel directory in Keras 3.
# Wrap the exported graph as an inference-only layer:
layer = keras.layers.TFSMLayer('saved_model_dir/', call_endpoint='serving_default')
# or go low-level (signatures only, no Keras object):
loaded = tf.saved_model.load('saved_model_dir/')

# Weights only (model architecture must already exist)
model = build_model()
model.load_weights('weights.weights.h5')

Weights Only — Partial and Strict Loading

# Load all weights (strict — topology and shapes must match)
model.load_weights('weights.weights.h5')

# Skip layers whose weight shapes/counts don't match (transfer learning)
model.load_weights('pretrained.weights.h5', skip_mismatch=True)

# Restore specific layers between differently-shaped models: copy per layer
for src, dst in zip(base_model.layers, model.layers):
    dst.set_weights(src.get_weights())

by_name=True was a legacy-HDF5 argument removed in Keras 3 — load_weights matches by layer topology; use skip_mismatch=True or per-layer set_weights() for partial restores.

TF Checkpoints

tf.train.Checkpoint tracks tf.Variable objects and optimizer state independently of Keras.

ckpt = tf.train.Checkpoint(
    model=model,
    optimizer=optimizer,
    step=tf.Variable(0, dtype=tf.int64),
)
manager = tf.train.CheckpointManager(ckpt, './ckpts', max_to_keep=3)

# Save
ckpt.step.assign_add(1)
manager.save()             # saves to ./ckpts/ckpt-1

# Restore latest
ckpt.restore(manager.latest_checkpoint)
if manager.latest_checkpoint:
    print(f'Restored from {manager.latest_checkpoint}')
else:
    print('Starting from scratch.')

# Restore specific checkpoint
ckpt.restore('./ckpts/ckpt-5')

# expect_partial: silences warnings when not all variables are restored
ckpt.restore(manager.latest_checkpoint).expect_partial()

# assert_consumed: errors if any saved variable was not restored
ckpt.restore(manager.latest_checkpoint).assert_consumed()

Custom Objects and Serialization

# Register for automatic deserialization (Keras 3)
@keras.saving.register_keras_serializable(package='MyPkg')
class MyLayer(keras.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def call(self, x):
        return x * self.units

    def get_config(self):
        config = super().get_config()
        config['units'] = self.units
        return config

# Load with custom objects dict (if not registered)
model = keras.models.load_model(
    'model.keras',
    custom_objects={'MyLayer': MyLayer},
)

# Globally register custom objects
with keras.saving.custom_object_scope({'MyLayer': MyLayer}):
    model = keras.models.load_model('model.keras')

Export for Serving

TF Serving

# Export as SavedModel (default for TF Serving)
model.export('saved_model_dir/')   # Keras 3 preferred
# or
tf.saved_model.save(model, 'saved_model_dir/')

# Inspect exported signatures
loaded = tf.saved_model.load('saved_model_dir/')
print(list(loaded.signatures.keys()))   # ['serving_default']
infer = loaded.signatures['serving_default']
result = infer(input_1=tf.constant(x))

TensorFlow Lite

# Convert to TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

# Post-training quantization
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# Float16 quantization
converter.target_spec.supported_types = [tf.float16]
# INT8 full quantization (requires representative dataset)
def representative_dataset():
    for x, _ in val_ds.take(100):
        yield [x]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type  = tf.int8
converter.inference_output_type = tf.int8
tflite_quant = converter.convert()

# Run TFLite model
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
in_idx  = interpreter.get_input_details()[0]['index']
out_idx = interpreter.get_output_details()[0]['index']
interpreter.set_tensor(in_idx, x_sample)
interpreter.invoke()
output = interpreter.get_tensor(out_idx)

TensorFlow.js

# pip install tensorflowjs
import tensorflowjs as tfjs

tfjs.converters.save_keras_model(model, 'tfjs_model/')
# or from SavedModel:
# tensorflowjs_converter --input_format=tf_saved_model saved_model_dir/ tfjs_model/

Saving Architecture Only

# JSON
json_config = model.to_json()
with open('arch.json', 'w') as f:
    f.write(json_config)

model_from_json = keras.models.model_from_json(json_config)

# YAML (Keras 2 only)
yaml_config = model.to_yaml()
model = keras.models.model_from_yaml(yaml_config)

Callbacks-Based Saving (Best Practice)

model.fit(
    train_ds,
    epochs=50,
    callbacks=[
        keras.callbacks.ModelCheckpoint(
            filepath='checkpoints/model_{epoch:03d}.keras',
            save_best_only=True,
            monitor='val_loss',
        ),
        keras.callbacks.BackupAndRestore(backup_dir='/tmp/backup'),
    ],
)

# Resume from last checkpoint
model = keras.models.load_model('checkpoints/model_042.keras')
model.fit(train_ds, epochs=100, initial_epoch=42, ...)

Common Gotchas

  • model.save() on a subclassed model without get_config saves the weights but not the architecture — you must reconstruct the model class yourself before calling load_weights.
  • Lambda layers do not serialize; replace with a proper Layer subclass.
  • initial_epoch must be passed to fit() when resuming so metrics and callbacks report correctly.
  • Custom loss/metric classes need get_config + @register_keras_serializable for load_model to work without passing custom_objects.
  • SavedModel vs .keras: SavedModel embeds the computation graph (useful for serving); .keras stores pure Python config + weights and requires Keras at load time.