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
| Format | Extension | Saves | Best for |
|---|---|---|---|
| Keras native | .keras | architecture + weights + optimizer + config | Keras 3 default; fully reproducible |
| SavedModel | directory | architecture + weights + computation graph | TF Serving, TFLite, TF.js export |
| HDF5 | .h5 | architecture + weights + optimizer | Legacy; avoid for new projects |
| Weights only | .weights.h5 (extension required) | weights tensors only | Partial 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=Truewas a legacy-HDF5 argument removed in Keras 3 —load_weightsmatches by layer topology; useskip_mismatch=Trueor per-layerset_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 withoutget_configsaves the weights but not the architecture — you must reconstruct the model class yourself before callingload_weights.- Lambda layers do not serialize; replace with a proper
Layersubclass. initial_epochmust be passed tofit()when resuming so metrics and callbacks report correctly.- Custom loss/metric classes need
get_config+@register_keras_serializableforload_modelto work without passingcustom_objects. - SavedModel vs
.keras: SavedModel embeds the computation graph (useful for serving);.kerasstores pure Python config + weights and requires Keras at load time.