tessl/pypi-tensorflow
An end-to-end open source platform for machine learning
—
Pending
saved-model.mddocs/
Model Management
Complete model serialization, checkpointing, and deployment utilities for production and inference. These operations provide comprehensive model lifecycle management capabilities.
Capabilities
Model Saving and Loading
Save and load complete models with all weights, architecture, and training configuration.
def save(obj, export_dir, signatures=None, options=None):
"""
Exports a tf.Module (and subclasses) obj to SavedModel format.
Parameters:
- obj: A trackable object (e.g. tf.Module or tf.keras.Model) to export
- export_dir: A directory in which to write the SavedModel
- signatures: Optional, either a tf.function with an input signature specified or a dictionary
- options: Optional, tf.saved_model.SaveOptions object that specifies options for saving
"""
def load(export_dir, tags=None, options=None):
"""
Load a SavedModel from export_dir.
Parameters:
- export_dir: The SavedModel directory to load from
- tags: A tag or sequence of tags identifying the MetaGraph to load
- options: Optional, tf.saved_model.LoadOptions object that specifies options for loading
Returns:
A trackable object with a save method
"""
def contains_saved_model(export_dir):
"""
Checks whether the provided export directory could contain a SavedModel.
Parameters:
- export_dir: Absolute or relative path to a directory containing the SavedModel
Returns:
True if the export directory contains SavedModel files, False otherwise
"""Checkpointing
Save and restore model weights and training state for resuming training.
class Checkpoint:
"""
Groups trackable objects, saving and restoring them.
Methods:
- save(file_prefix): Saves a training checkpoint and provides a context manager
- restore(save_path): Restore a training checkpoint
- read(save_path): Returns CheckpointReader for checkpoint inspection
"""
def __init__(self, **kwargs):
"""
Groups trackable objects, saving and restoring them.
Parameters:
- **kwargs: Keyword arguments are set as attributes of this object, and are saved with the checkpoint
"""
def save(self, file_prefix, session=None):
"""
Saves a training checkpoint and provides a context manager.
Parameters:
- file_prefix: A prefix to use for the checkpoint filenames
- session: The session to evaluate variables in. Ignored when executing eagerly
Returns:
The full path to the checkpoint
"""
def restore(self, save_path):
"""
Restore a training checkpoint.
Parameters:
- save_path: The path to the checkpoint, as returned by save or tf.train.latest_checkpoint
Returns:
A load status object, which can be used to make assertions about the status of a checkpoint restoration
"""
def read(self, save_path):
"""
Returns a CheckpointReader for the checkpoint.
Parameters:
- save_path: The path to the checkpoint, as returned by save or tf.train.latest_checkpoint
Returns:
A CheckpointReader object
"""
class CheckpointManager:
"""
Deletes old checkpoints.
Methods:
- save(checkpoint_number): Creates a new checkpoint
"""
def __init__(self, checkpoint, directory, max_to_keep=5, keep_checkpoint_every_n_hours=None,
checkpoint_name="ckpt", step_counter=None, checkpoint_interval=None,
init_fn=None):
"""
Deletes old checkpoints.
Parameters:
- checkpoint: The tf.train.Checkpoint instance to save and manage checkpoints for
- directory: The path to a directory in which to write checkpoints
- max_to_keep: An integer, the number of checkpoints to keep
- keep_checkpoint_every_n_hours: Upon removal, keep checkpoints every N hours
- checkpoint_name: Custom name for the checkpoint file
- step_counter: A tf.Variable instance for checking the current step counter value
- checkpoint_interval: An integer, indicates that keep_checkpoint_every_n_hours should be based on checkpoints saved every checkpoint_interval steps
- init_fn: Callable. Function executed the first time a checkpoint is saved
"""
def save(self, checkpoint_number=None, check_interval=True):
"""
Creates a new checkpoint and manages deletion of old checkpoints.
Parameters:
- checkpoint_number: An optional integer, or an integer-dtype Variable or Tensor, used to number the checkpoint
- check_interval: An optional boolean. The default behaviour is that checkpoint_interval is ignored when checkpoint_number is provided
Returns:
The path to the new checkpoint. It is also recorded in the checkpoints and latest_checkpoint properties
"""Checkpoint Utilities
Utility functions for working with checkpoints.
def list_variables(checkpoint_dir):
"""
Returns list of all variables in the checkpoint.
Parameters:
- checkpoint_dir: Directory with checkpoint file or path to checkpoint
Returns:
List of tuples (name, shape) for all variables in the checkpoint
"""
def load_checkpoint(checkpoint_dir):
"""
Returns CheckpointReader for checkpoint found in checkpoint_dir.
Parameters:
- checkpoint_dir: Directory with checkpoint file or path to checkpoint
Returns:
CheckpointReader instance
"""
def load_variable(checkpoint_dir, name):
"""
Returns the tensor value of the given variable in the checkpoint.
Parameters:
- checkpoint_dir: Directory with checkpoint file or path to checkpoint
- name: Name of the variable to return
Returns:
A numpy ndarray with a copy of the value of this variable
"""
def latest_checkpoint(checkpoint_dir, latest_filename=None):
"""
Finds the filename of latest saved checkpoint file.
Parameters:
- checkpoint_dir: Directory where the variables were saved
- latest_filename: Optional name for the protocol buffer file that contains the list of most recent checkpoint filenames
Returns:
The full path to the latest checkpoint or None if no checkpoint was found
"""SavedModel Utilities
Additional utilities for working with SavedModel format.
class SaveOptions:
"""
Options for saving to SavedModel.
Parameters:
- namespace_whitelist: List of strings containing op namespaces to whitelist when saving a model
- save_debug_info: Boolean indicating whether debug information is saved
- function_aliases: Optional dictionary of string -> string of function aliases
- experimental_io_device: string. Applies in a distributed setting
- experimental_variable_policy: The policy to apply to variables when saving
"""
class LoadOptions:
"""
Options for loading a SavedModel.
Parameters:
- allow_partial_checkpoint: Boolean. Defaults to False. When enabled, allows the SavedModel checkpoint to be missing variables
- experimental_io_device: string. Loads SavedModel and variables on the specified device
- experimental_skip_checkpoint: boolean. If True, the checkpoint will not be loaded, and the SavedModel will be loaded with randomly initialized variable values
"""
class Asset:
"""
Represents a file asset to copy into the SavedModel.
Parameters:
- path: A path, or a 0-D tf.string Tensor with path to the asset
"""Usage Examples
import tensorflow as tf
import os
# Create a simple model
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
# Save entire model to SavedModel format
tf.saved_model.save(model, 'my_saved_model')
# Load the saved model
loaded_model = tf.saved_model.load('my_saved_model')
# For Keras models, use keras save/load for full functionality
model.save('my_keras_model.h5')
loaded_keras_model = tf.keras.models.load_model('my_keras_model.h5')
# Checkpoint example
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
# Create checkpoint object
checkpoint = tf.train.Checkpoint(optimizer=tf.keras.optimizers.Adam(),
model=model)
# Save checkpoint
checkpoint.save(file_prefix=checkpoint_prefix)
# Restore from checkpoint
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
# Using CheckpointManager for automatic cleanup
manager = tf.train.CheckpointManager(
checkpoint, directory=checkpoint_dir, max_to_keep=3
)
# Save with automatic cleanup
save_path = manager.save()
print(f"Saved checkpoint for step {step}: {save_path}")
# Training loop with checkpointing
optimizer = tf.keras.optimizers.Adam()
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
manager = tf.train.CheckpointManager(checkpoint, './checkpoints', max_to_keep=3)
# Restore if checkpoint exists
checkpoint.restore(manager.latest_checkpoint)
if manager.latest_checkpoint:
print(f"Restored from {manager.latest_checkpoint}")
else:
print("Initializing from scratch.")
# Training step function
@tf.function
def train_step(x, y):
with tf.GradientTape() as tape:
predictions = model(x, training=True)
loss = tf.keras.losses.binary_crossentropy(y, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return loss
# Training loop
for epoch in range(10):
# Training code here...
# x_batch, y_batch = get_batch()
# loss = train_step(x_batch, y_batch)
# Save checkpoint every few epochs
if epoch % 2 == 0:
save_path = manager.save()
print(f"Saved checkpoint for epoch {epoch}: {save_path}")
# Inspect checkpoint contents
checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
if checkpoint_path:
variables = tf.train.list_variables(checkpoint_path)
for name, shape in variables:
print(f"Variable: {name}, Shape: {shape}")
# Load specific variable
specific_var = tf.train.load_variable(checkpoint_path, 'model/dense/kernel/.ATTRIBUTES/VARIABLE_VALUE')
print(f"Loaded variable shape: {specific_var.shape}")
# Check if directory contains SavedModel
if tf.saved_model.contains_saved_model('my_saved_model'):
print("Directory contains a valid SavedModel")
# Advanced SavedModel with custom signatures
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
def inference_func(x):
return model(x)
# Save with custom signature
tf.saved_model.save(
model,
'model_with_signature',
signatures={'serving_default': inference_func}
)Install with Tessl CLI
npx tessl i tessl/pypi-tensorflow