ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

Migrate the SavedModel workflow

View on TensorFlow.org Run in Google Colab View source on GitHub Download notebook

Once you have migrated your model from TensorFlow 1's graphs and sessions to TensorFlow 2 APIs, such as tf.function, tf.Module, and tf.keras.Model, you can migrate the model saving and loading code. This notebook provides examples of how you can save and load in the SavedModel format in TensorFlow 1 and TensorFlow 2. Here is a quick overview of the related API changes for migration from TensorFlow 1 to TensorFlow 2:

TensorFlow 1 Migration to TensorFlow 2
Saving tf.compat.v1.saved_model.Builder
tf.compat.v1.saved_model.simple_save
tf.saved_model.save
Keras: tf.keras.models.save_model
Loading tf.compat.v1.saved_model.load tf.saved_model.load
Keras: tf.keras.models.load_model
Signatures: a set of input
and output tensors that
can be used to run the
Generated using the *.signature_def utils
(e.g. tf.compat.v1.saved_model.predict_signature_def)
Write a tf.function and export it using the signatures argument
in tf.saved_model.save.
Classification
and regression
:
special types of signatures
Generated with
tf.compat.v1.saved_model.classification_signature_def,
tf.compat.v1.saved_model.regression_signature_def,
and certain Estimator exports.
These two signature types have been removed from TensorFlow 2.
If the serving library requires these method names,
tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater.

For a more in-depth explanation of the mapping, refer to the Changes from TensorFlow 1 to TensorFlow 2 section below.

Setup

The examples below show how to export and load the same dummy TensorFlow model (defined as add_two below) to a SavedModel format using the TensorFlow 1 and TensorFlow 2 APIs. Start by setting up the imports and utility functions:

import tensorflow as tf
import tensorflow.compat.v1 as tf1
import shutil

def remove_dir(path):
  try:
    shutil.rmtree(path)
  except:
    pass

def add_two(input):
  return input + 2

TensorFlow 1: Save and export a SavedModel

In TensorFlow 1, you use the tf.compat.v1.saved_model.Builder, tf.compat.v1.saved_model.simple_save, and tf.estimator.Estimator.export_saved_model APIs to build, save, and export the TensorFlow graph and session:

1. Save the graph as a SavedModel with SavedModelBuilder

remove_dir("saved-model-builder")

with tf.Graph().as_default() as g:
  with tf1.Session() as sess:
    input = tf1.placeholder(tf.float32, shape=[])
    output = add_two(input)
    print("add two output: ", sess.run(output, {input: 3.}))

    # Save with SavedModelBuilder
    builder = tf1.saved_model.Builder('saved-model-builder')
    sig_def = tf1.saved_model.predict_signature_def(
        inputs={'input': input},
        outputs={'output': output})
    builder.add_meta_graph_and_variables(
        sess, tags=["serve"], signature_def_map={
            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: sig_def
    })
    builder.save()
add two output:  5.0
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/saved_model/signature_def_utils_impl.py:201: build_tensor_info (from tensorflow.python.saved_model.utils_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.utils.build_tensor_info or tf.compat.v1.saved_model.build_tensor_info.
INFO:tensorflow:No assets to save.
INFO:tensorflow:No assets to write.
INFO:tensorflow:SavedModel written to: saved-model-builder/saved_model.pb
!saved_model_cli run --dir simple-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10
Traceback (most recent call last):
  File "/tmpfs/src/tf_docs_env/bin/saved_model_cli", line 8, in <module>
    sys.exit(main())
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py", line 1192, in main
    args.func(args)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py", line 752, in run
    init_tpu=args.init_tpu, tf_debug=args.tf_debug)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py", line 409, in run_saved_model_with_feed_dict
    tag_set)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_utils.py", line 117, in get_meta_graph_def
    saved_model = read_saved_model(saved_model_dir)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_utils.py", line 55, in read_saved_model
    raise IOError("SavedModel file does not exist at: %s" % saved_model_dir)
OSError: SavedModel file does not exist at: simple-save

2. Build a SavedModel for serving

remove_dir("simple-save")

with tf.Graph().as_default() as g:
  with tf1.Session() as sess:
    input = tf1.placeholder(tf.float32, shape=[])
    output = add_two(input)
    print("add_two output: ", sess.run(output, {input: 3.}))

    tf1.saved_model.simple_save(
        sess, 'simple-save',
        inputs={'input': input},
        outputs={'output': output})
add_two output:  5.0
WARNING:tensorflow:From /tmp/ipykernel_14881/250978412.py:12: simple_save (from tensorflow.python.saved_model.simple_save) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.simple_save.
INFO:tensorflow:Assets added to graph.
INFO:tensorflow:No assets to write.
INFO:tensorflow:SavedModel written to: simple-save/saved_model.pb
!saved_model_cli run --dir simple-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py:445: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.
INFO:tensorflow:Saver not created because there are no variables in the graph to restore
INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.
Result for output key output:
12.0

3. Export the Estimator inference graph as a SavedModel

In the definition of the Estimator model_fn (defined below), you can define signatures in your model by returning export_outputs in the tf.estimator.EstimatorSpec. There are different types of outputs:

These will produce classification, regression, and prediction signature types, respectively.

When the estimator is exported with tf.estimator.Estimator.export_saved_model, these signatures will be saved with the model.

def model_fn(features, labels, mode):
  output = add_two(features['input'])
  step = tf1.train.get_global_step()
  return tf.estimator.EstimatorSpec(
      mode,
      predictions=output,
      train_op=step.assign_add(1),
      loss=tf.constant(0.),
      export_outputs={
          tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: \
          tf.estimator.export.PredictOutput({'output': output})})
est = tf.estimator.Estimator(model_fn, 'estimator-checkpoints')

# Train for one step to create a checkpoint.
def train_fn():
  return tf.data.Dataset.from_tensors({'input': 3.})
est.train(train_fn, steps=1)

# This utility function `build_raw_serving_input_receiver_fn` takes in raw
# tensor features and builds an "input serving receiver function", which
# creates placeholder inputs to the model.
serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(
    {'input': tf.constant(3.)})  # Pass in a dummy input batch.
estimator_path = est.export_saved_model('exported-estimator', serving_input_fn)

# Estimator's export_saved_model creates a time stamped directory. Move this
# to a set path so it can be inspected with `saved_model_cli` in the cell below.
!rm -rf estimator-model
import shutil
shutil.move(estimator_path, 'estimator-model')
INFO:tensorflow:Using default config.
INFO:tensorflow:Using config: {'_model_dir': 'estimator-checkpoints', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true
graph_options {
  rewrite_options {
    meta_optimizer_iterations: ONE
  }
}
, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_checkpoint_save_graph_def': True, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/training/training_util.py:236: Variable.initialized_value (from tensorflow.python.ops.variables) is deprecated and will be removed in a future version.
Instructions for updating:
Use Variable.read_value. Variables in 2.X are initialized automatically both in eager and graph (inside tf.defun) contexts.
INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Create CheckpointSaverHook.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
INFO:tensorflow:Done running local_init_op.
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0...
INFO:tensorflow:Saving checkpoints for 0 into estimator-checkpoints/model.ckpt.
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0...
INFO:tensorflow:loss = 0.0, step = 1
INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 1...
INFO:tensorflow:Saving checkpoints for 1 into estimator-checkpoints/model.ckpt.
INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 1...
INFO:tensorflow:Loss for final step: 0.0.
INFO:tensorflow:Calling model_fn.
INFO:tensorflow:Done calling model_fn.
INFO:tensorflow:Signatures INCLUDED in export for Classify: None
INFO:tensorflow:Signatures INCLUDED in export for Regress: None
INFO:tensorflow:Signatures INCLUDED in export for Predict: ['serving_default']
INFO:tensorflow:Signatures INCLUDED in export for Train: None
INFO:tensorflow:Signatures INCLUDED in export for Eval: None
INFO:tensorflow:Restoring parameters from estimator-checkpoints/model.ckpt-1
INFO:tensorflow:Assets added to graph.
INFO:tensorflow:No assets to write.
INFO:tensorflow:SavedModel written to: exported-estimator/temp-1632341678/saved_model.pb
'estimator-model'
!saved_model_cli run --dir estimator-model --tag_set serve \
 --signature_def serving_default --input_exprs input=[10]
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py:445: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.
INFO:tensorflow:Restoring parameters from estimator-model/variables/variables
Result for output key output:
[12.]

TensorFlow 2: Save and export a SavedModel

Save and export a SavedModel defined with tf.Module

To export your model in TensorFlow 2, you must define a tf.Module or a tf.keras.Model to hold all of your model's variables and functions. Then, you can call tf.saved_model.save to create a SavedModel. Refer to Saving a custom model in the Using the SavedModel format guide to learn more.

class MyModel(tf.Module):
  @tf.function
  def __call__(self, input):
    return add_two(input)

model = MyModel()

@tf.function
def serving_default(input):
  return {'output': model(input)}

signature_function = serving_default.get_concrete_function(
    tf.TensorSpec(shape=[], dtype=tf.float32))
tf.saved_model.save(
    model, 'tf2-save', signatures={
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_function})
INFO:tensorflow:Assets written to: tf2-save/assets
2021-09-22 20:14:41.714399: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
!saved_model_cli run --dir tf2-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py:445: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.
INFO:tensorflow:Restoring parameters from tf2-save/variables/variables
Result for output key output:
12.0

Save and export a SavedModel defined with Keras

The Keras APIs for saving and exporting—Mode.save or tf.keras.models.save_model—can export a SavedModel from a tf.keras.Model. Check out the Save and load Keras models for more details.

inp = tf.keras.Input(3)
out = add_two(inp)
model = tf.keras.Model(inputs=inp, outputs=out)

@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
def serving_default(input):
  return {'output': model(input)}

model.save('keras-model', save_format='tf', signatures={
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: serving_default})
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Model was constructed with shape (None, 3) for input KerasTensor(type_spec=TensorSpec(shape=(None, 3), dtype=tf.float32, name='input_1'), name='input_1', description="created by layer 'input_1'"), but it was called on an input with incompatible shape ().
INFO:tensorflow:Assets written to: keras-model/assets
!saved_model_cli run --dir keras-model --tag_set serve \
 --signature_def serving_default --input_exprs input=10
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/tools/saved_model_cli.py:445: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.
INFO:tensorflow:Restoring parameters from keras-model/variables/variables
Result for output key output:
12.0

Loading a SavedModel

A SavedModel saved with any of the above APIs can be loaded using either TensorFlow 1 or TensorFlow APIs.

A TensorFlow 1 SavedModel can generally be used for inference when loaded into TensorFlow 2, but training (generating gradients) is only possible if the SavedModel contains resource variables. You can check the dtype of the variables—if the variable dtype contains "_ref", then it is a reference variable.

A TensorFlow 2 SavedModel can be loaded and executed from TensorFlow 1 as long as the SavedModel is saved with signatures.

The sections below contain code samples showing how to load the SavedModels saved in the previous sections, and call the exported signature.

TensorFlow 1: Load a SavedModel with tf.saved_model.load

In TensorFlow 1, you can import a SavedModel directly into the current graph and session using tf.saved_model.load. You can call Session.run on the tensor input and output names:

def load_tf1(path, input):
  print('Loading from', path)
  with tf.Graph().as_default() as g:
    with tf1.Session() as sess:
      meta_graph = tf1.saved_model.load(sess, ["serve"], path)
      sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
      input_name = sig_def.inputs['input'].name
      output_name = sig_def.outputs['output'].name
      print('  Output with input', input, ': ', 
            sess.run(output_name, feed_dict={input_name: input}))

load_tf1('saved-model-builder', 5.)
load_tf1('simple-save', 5.)
load_tf1('estimator-model', [5.])  # Estimator's input must be batched.
load_tf1('tf2-save', 5.)
load_tf1('keras-model', 5.)
Loading from saved-model-builder
WARNING:tensorflow:From /tmp/ipykernel_14881/1548963983.py:5: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.
INFO:tensorflow:Saver not created because there are no variables in the graph to restore
INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.
  Output with input 5.0 :  7.0
Loading from simple-save
INFO:tensorflow:Saver not created because there are no variables in the graph to restore
INFO:tensorflow:The specified SavedModel has no variables; no checkpoints were restored.
  Output with input 5.0 :  7.0
Loading from estimator-model
INFO:tensorflow:Restoring parameters from estimator-model/variables/variables
  Output with input [5.0] :  [7.]
Loading from tf2-save
INFO:tensorflow:Restoring parameters from tf2-save/variables/variables
  Output with input 5.0 :  7.0
Loading from keras-model
INFO:tensorflow:Restoring parameters from keras-model/variables/variables
  Output with input 5.0 :  7.0

TensorFlow 2: Load a model saved with tf.saved_model

In TensorFlow 2, objects are loaded into a Python object that stores the variables and functions. This is compatible with models saved from TensorFlow 1.

Check out the tf.saved_model.load API docs and Loading and using a custom model section from the Using the SavedModel format guide for details.

def load_tf2(path, input):
  print('Loading from', path)
  loaded = tf.saved_model.load(path)
  out = loaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY](
      tf.constant(input))['output']
  print('  Output with input', input, ': ', out)

load_tf2('saved-model-builder', 5.)
load_tf2('simple-save', 5.)
load_tf2('estimator-model', [5.])  # Estimator's input must be batched.
load_tf2('tf2-save', 5.)
load_tf2('keras-model', 5.)
Loading from saved-model-builder
INFO:tensorflow:Saver not created because there are no variables in the graph to restore
  Output with input 5.0 :  tf.Tensor(7.0, shape=(), dtype=float32)
Loading from simple-save
INFO:tensorflow:Saver not created because there are no variables in the graph to restore
  Output with input 5.0 :  tf.Tensor(7.0, shape=(), dtype=float32)
Loading from estimator-model
  Output with input [5.0] :  tf.Tensor([7.], shape=(1,), dtype=float32)
Loading from tf2-save
  Output with input 5.0 :  tf.Tensor(7.0, shape=(), dtype=float32)
Loading from keras-model
  Output with input 5.0 :  tf.Tensor(7.0, shape=(), dtype=float32)

Models saved with the TensorFlow 2 API can also access tf.functions and variables that are attached to the model (instead of those exported as signatures). For example:

loaded = tf.saved_model.load('tf2-save')
print('restored __call__:', loaded.__call__)
print('output with input 5.', loaded(5))
restored __call__: <tensorflow.python.saved_model.function_deserialization.RestoredFunction object at 0x7f0714ed3f50>
output with input 5. tf.Tensor(7.0, shape=(), dtype=float32)

TensorFlow 2: Load a model saved with Keras

The Keras loading API—tf.keras.models.load_model—allows you to reload a saved model back into a Keras Model object. Note that this only allows you to load SavedModels saved with Keras (Model.save or tf.keras.models.save_model).

Models saved with tf.saved_model.save should be loaded with tf.saved_model.load. You can load a Keras model saved with Model.save using tf.saved_model.load but you will only get the TensorFlow graph. Refer to the tf.keras.models.load_model API docs and Save and load Keras models guide for details.

loaded_model = tf.keras.models.load_model('keras-model')
loaded_model.predict_on_batch(tf.constant([1, 3, 4]))
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:tensorflow:Model was constructed with shape (None, 3) for input KerasTensor(type_spec=TensorSpec(shape=(None, 3), dtype=tf.float32, name='input_1'), name='input_1', description="created by layer 'input_1'"), but it was called on an input with incompatible shape (3, 1).
array([[3.],
       [5.],
       [6.]], dtype=float32)

GraphDef and MetaGraphDef

There is no straightforward way to load a raw GraphDef or MetaGraphDef to TF2. However, you can convert the TF1 code that imports the graph into a TF2 concrete_function using v1.wrap_function.

First, save a MetaGraphDef:

# Save a simple multiplication computation:
with tf.Graph().as_default() as g:
  x = tf1.placeholder(tf.float32, shape=[], name='x')
  v = tf.Variable(3.0, name='v')
  y = tf.multiply(x, v, name='y')
  with tf1.Session() as sess:
    sess.run(v.initializer)
    print(sess.run(y, feed_dict={x: 5}))
    s = tf1.train.Saver()
    s.export_meta_graph('multiply.pb', as_text=True)
    s.save(sess, 'multiply_values.ckpt')
15.0

Using TF1 APIs, you can use tf1.train.import_meta_graph to import the graph and restore the values:

with tf.Graph().as_default() as g:
  meta = tf1.train.import_meta_graph('multiply.pb')
  x = g.get_tensor_by_name('x:0')
  y = g.get_tensor_by_name('y:0')
  with tf1.Session() as sess:
    meta.restore(sess, 'multiply_values.ckpt')
    print(sess.run(y, feed_dict={x: 5}))
INFO:tensorflow:Restoring parameters from multiply_values.ckpt
15.0

There are no TF2 APIs for loading the graph, but you can still import it into a concrete function that can be executed in eager mode:

def import_multiply():
  # Any graph-building code is allowed here.
  tf1.train.import_meta_graph('multiply.pb')

# Creates a tf.function with all the imported elements in the function graph.
wrapped_import = tf1.wrap_function(import_multiply, [])
import_graph = wrapped_import.graph
x = import_graph.get_tensor_by_name('x:0')
y = import_graph.get_tensor_by_name('y:0')

# Restore the variable values.
tf1.train.Saver(wrapped_import.variables).restore(
    sess=None, save_path='multiply_values.ckpt')

# Create a concrete function by pruning the wrap_function (similar to sess.run).
multiply_fn = wrapped_import.prune(feeds=x, fetches=y)

# Run this function
multiply_fn(tf.constant(5.))  # inputs to concrete functions must be Tensors.
WARNING:tensorflow:Saver is deprecated, please switch to tf.train.Checkpoint or tf.keras.Model.save_weights for training checkpoints. When executing eagerly variables do not necessarily have unique names, and so the variable.name-based lookups Saver performs are error-prone.
INFO:tensorflow:Restoring parameters from multiply_values.ckpt
<tf.Tensor: shape=(), dtype=float32, numpy=15.0>

Changes from TensorFlow 1 to TensorFlow 2

This section lists out key saving and loading terms from TensorFlow 1, their TensorFlow 2 equivalents, and what has changed.

SavedModel

SavedModel is a format that stores a complete TensorFlow program with parameters and computation. It contains signatures used by serving platforms to run the model.

The file format itself has not changed significantly, so SavedModels can be loaded and served using either TensorFlow 1 or TensorFlow 2 APIs.

Differences between TensorFlow 1 and TensorFlow 2

The serving and inference use cases have not been updated in TensorFlow 2, aside from the API changes—the improvement was introduced in the ability to reuse and compose models loaded from SavedModel.

In TensorFlow 2, the program is represented by objects like tf.Variable, tf.Module, or higher-level Keras models (tf.keras.Model) and layers (tf.keras.layers). There are no more global variables that have values stored in a session, and the graph now exists in different tf.functions. Consequently, during a model export, SavedModel saves each component and function graphs separately.

When you write a TensorFlow program with the TensorFlow Python APIs, you must build an object to manage the variables, functions, and other resources. Generally, this is accomplished by using the Keras APIs, but you can also build the object by creating or subclassing tf.Module.

Keras models (tf.keras.Model) and tf.Module automatically track variables and functions attached to them. SavedModel saves these connections between modules, variables, and functions, so that they can be restored when loading.

Signatures

Signatures are the endpoints of a SavedModel—they tell the user how to run the model and what inputs are needed.

In TensorFlow 1, signatures are created by listing the input and output tensors. In TensorFlow 2, signatures are generated by passing in concrete functions. (Read more about TensorFlow functions in the Introduction to graphs and tf.function guide.) In short, a concrete function is generated from a tf.function:

# Option 1: Specify an input signature.
@tf.function(input_signature=[...])
def fn(...):
  ...
  return outputs

tf.saved_model.save(model, path, signatures={
    'name': fn
})
# Option 2: Call `get_concrete_function`
@tf.function
def fn(...):
  ...
  return outputs

tf.saved_model.save(model, path, signatures={
    'name': fn.get_concrete_function(...)
})

Session.run

In TensorFlow 1, you could call Session.run with the imported graph as long as you already know the tensor names. This allows you to retrieve the restored variable values, or run parts of the model that were not exported in the signatures.

In TensorFlow 2, you can directly access a variable, such as a weights matrix (kernel):

model = tf.Module()
model.dense_layer = tf.keras.layers.Dense(...)
tf.saved_model.save('my_saved_model')
loaded = tf.saved_model.load('my_saved_model')
loaded.dense_layer.kernel

or call tf.functions attached to the model object: for example, loaded.__call__.

Unlike TF1, there is no way to extract parts of a function and access intermediate values. You must export all of the needed functionality in the saved object.

TensorFlow Serving migration notes

SavedModel was originally created to work with TensorFlow Serving. This platform offers different types of prediction requests: classify, regress, and predict.

The TensorFlow 1 API allows you to create these types of signatures with the utils:

Classification (classification_signature_def) and regression (regression_signature_def) restrict the inputs and outputs, so the inputs must be a tf.Example, and the outputs must be classes, scores or prediction. Meanwhile, the predict signature (predict_signature_def) has no restrictions.

SavedModels exported with the TensorFlow 2 API are compatible with TensorFlow Serving, but will only contain prediction signatures. The classification and regression signatures have been removed.

If you require the use of the classification and regression signatures, you may modify the exported SavedModel using tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater.

Next steps

To learn more about SavedModels in TensorFlow 2, check out the following guides:

If you are using TensorFlow Hub, you may find these guides useful: