TensorFlow Cheatsheet

Saving and Loading

Use this TensorFlow reference while you build software engineering projects, review code, or refresh the syntax you reach for most.

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.