<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2018 The TensorFlow Authors.
Step1:  Using the SavedModel format
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  You'll use an image of Grace Hopper as a running example, and a Keras pre-trained image classification model since it's easy to use. Custom models work too, and are covered in detail later.
Step3:  The top prediction for this image is "military uniform".
Step4:  The save-path follows a convention used by TensorFlow Serving where the last path component (1/ here) is a version number for your model - it allows tools like Tensorflow Serving to reason about the relative freshness.
You can load the SavedModel back into Python with tf.saved_model.load and see how Admiral Hopper's image is classified.
Step5:  Imported signatures always return dictionaries. To customize signature names and output dictionary keys, see Specifying signatures during export.
Step6:  Running inference from the SavedModel gives the same result as the original model.
Step7:  Running a SavedModel in TensorFlow Serving
SavedModels are usable from Python (more on that below), but production environments typically use a dedicated service for inference without running Python code. This is easy to set up from a SavedModel using TensorFlow Serving.
See the TensorFlow Serving REST tutorial for an end-to-end tensorflow-serving example.
The SavedModel format on disk
A SavedModel is a directory containing serialized signatures and the state needed to run them, including variable values and vocabularies.
Step8:  The saved_model.pb file stores the actual TensorFlow program, or model, and a set of named signatures, each identifying a function that accepts tensor inputs and produces tensor outputs.
SavedModels may contain multiple variants of the model (multiple v1.MetaGraphDefs, identified with the --tag_set flag to saved_model_cli), but this is rare. APIs which create multiple variants of a model include tf.Estimator.experimental_export_all_saved_models and in TensorFlow 1.x tf.saved_model.Builder.
Step9:  The variables directory contains a standard training checkpoint (see the guide to training checkpoints).
Step10:  The assets directory contains files used by the TensorFlow graph, for example text files used to initialize vocabulary tables. It is unused in this example.
SavedModels may have an assets.extra directory for any files not used by the TensorFlow graph, for example information for consumers about what to do with the SavedModel. TensorFlow itself does not use this directory.
Saving a custom model
tf.saved_model.save supports saving tf.Module objects and its subclasses, like tf.keras.Layer and tf.keras.Model.
Let's look at an example of saving and restoring a tf.Module.
Step11:  When you save a tf.Module, any tf.Variable attributes, tf.function-decorated methods, and tf.Modules found via recursive traversal are saved. (See the Checkpoint tutorial for more about this recursive traversal.) However, any Python attributes, functions, and data are lost. This means that when a tf.function is saved, no Python code is saved.
If no Python code is saved, how does SavedModel know how to restore the function?
Briefly, tf.function works by tracing the Python code to generate a ConcreteFunction (a callable wrapper around tf.Graph). When saving a tf.function, you're really saving the tf.function's cache of ConcreteFunctions.
To learn more about the relationship between tf.function and ConcreteFunctions, refer to the tf.function guide.
Step12:  Loading and using a custom model
When you load a SavedModel in Python, all tf.Variable attributes, tf.function-decorated methods, and tf.Modules are restored in the same object structure as the original saved tf.Module.
Step13:  Because no Python code is saved, calling a tf.function with a new input signature will fail
Step14:  General fine-tuning
A SavedModel from Keras provides more details than a plain __call__ to address more advanced cases of fine-tuning. TensorFlow Hub recommends to provide the following of those, if applicable, in SavedModels shared for the purpose of fine-tuning
Step15:  Specifying signatures during export
Tools like TensorFlow Serving and saved_model_cli can interact with SavedModels. To help these tools determine which ConcreteFunctions to use, you need to specify serving signatures. tf.keras.Models automatically specify serving signatures, but you'll have to explicitly declare a serving signature for our custom modules.
IMPORTANT
Step16:  To declare a serving signature, specify a ConcreteFunction using the signatures kwarg. When specifying a single signature, its signature key will be 'serving_default', which is saved as the constant tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY.
Step17:  To export multiple signatures, pass a dictionary of signature keys to ConcreteFunctions. Each signature key corresponds to one ConcreteFunction.
Step18:  By default, the output tensor names are fairly generic, like output_0. To control the names of outputs, modify your tf.function to return a dictionary that maps output names to outputs. The names of inputs are derived from the Python function arg names.
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2018 The TensorFlow Authors.
End of explanation
import os
import tempfile
from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf
tmpdir = tempfile.mkdtemp()
physical_devices = tf.config.list_physical_devices('GPU')
for device in physical_devices:
  tf.config.experimental.set_memory_growth(device, True)
file = tf.keras.utils.get_file(
    "grace_hopper.jpg",
    "https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg")
img = tf.keras.utils.load_img(file, target_size=[224, 224])
plt.imshow(img)
plt.axis('off')
x = tf.keras.utils.img_to_array(img)
x = tf.keras.applications.mobilenet.preprocess_input(
    x[tf.newaxis,...])
Explanation: Using the SavedModel format
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/guide/saved_model"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/saved_model.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/docs/blob/master/site/en/guide/saved_model.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/saved_model.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>
A SavedModel contains a complete TensorFlow program, including trained parameters (i.e, tf.Variables) and computation. It does not require the original model building code to run, which makes it useful for sharing or deploying with TFLite, TensorFlow.js, TensorFlow Serving, or TensorFlow Hub.
You can save and load a model in the SavedModel format using the following APIs:
- Low-level tf.saved_model API. This document describes how to use this API in detail.
 - Save: tf.saved_model.save(model, path_to_dir)
 - Load: model = tf.saved_model.load(path_to_dir)
- High-level tf.keras.Model API. Refer to the keras save and serialize guide.
- If you just want to save/load weights during training, refer to the checkpoints guide.
Caution: TensorFlow models are code and it is important to be careful with untrusted code. Learn more in Using TensorFlow securely.
Creating a SavedModel from Keras
For a quick introduction, this section exports a pre-trained Keras model and serves image classification requests with it. The rest of the guide will fill in details and discuss other ways to create SavedModels.
End of explanation
labels_path = tf.keras.utils.get_file(
    'ImageNetLabels.txt',
    'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())
pretrained_model = tf.keras.applications.MobileNet()
result_before_save = pretrained_model(x)
decoded = imagenet_labels[np.argsort(result_before_save)[0,::-1][:5]+1]
print("Result before saving:\n", decoded)
Explanation: You'll use an image of Grace Hopper as a running example, and a Keras pre-trained image classification model since it's easy to use. Custom models work too, and are covered in detail later.
End of explanation
mobilenet_save_path = os.path.join(tmpdir, "mobilenet/1/")
tf.saved_model.save(pretrained_model, mobilenet_save_path)
Explanation: The top prediction for this image is "military uniform".
End of explanation
loaded = tf.saved_model.load(mobilenet_save_path)
print(list(loaded.signatures.keys()))  # ["serving_default"]
Explanation: The save-path follows a convention used by TensorFlow Serving where the last path component (1/ here) is a version number for your model - it allows tools like Tensorflow Serving to reason about the relative freshness.
You can load the SavedModel back into Python with tf.saved_model.load and see how Admiral Hopper's image is classified.
End of explanation
infer = loaded.signatures["serving_default"]
print(infer.structured_outputs)
Explanation: Imported signatures always return dictionaries. To customize signature names and output dictionary keys, see Specifying signatures during export.
End of explanation
labeling = infer(tf.constant(x))[pretrained_model.output_names[0]]
decoded = imagenet_labels[np.argsort(labeling)[0,::-1][:5]+1]
print("Result after saving and loading:\n", decoded)
Explanation: Running inference from the SavedModel gives the same result as the original model.
End of explanation
!ls {mobilenet_save_path}
Explanation: Running a SavedModel in TensorFlow Serving
SavedModels are usable from Python (more on that below), but production environments typically use a dedicated service for inference without running Python code. This is easy to set up from a SavedModel using TensorFlow Serving.
See the TensorFlow Serving REST tutorial for an end-to-end tensorflow-serving example.
The SavedModel format on disk
A SavedModel is a directory containing serialized signatures and the state needed to run them, including variable values and vocabularies.
End of explanation
!saved_model_cli show --dir {mobilenet_save_path} --tag_set serve
Explanation: The saved_model.pb file stores the actual TensorFlow program, or model, and a set of named signatures, each identifying a function that accepts tensor inputs and produces tensor outputs.
SavedModels may contain multiple variants of the model (multiple v1.MetaGraphDefs, identified with the --tag_set flag to saved_model_cli), but this is rare. APIs which create multiple variants of a model include tf.Estimator.experimental_export_all_saved_models and in TensorFlow 1.x tf.saved_model.Builder.
End of explanation
!ls {mobilenet_save_path}/variables
Explanation: The variables directory contains a standard training checkpoint (see the guide to training checkpoints).
End of explanation
class CustomModule(tf.Module):
  def __init__(self):
    super(CustomModule, self).__init__()
    self.v = tf.Variable(1.)
  @tf.function
  def __call__(self, x):
    print('Tracing with', x)
    return x * self.v
  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def mutate(self, new_v):
    self.v.assign(new_v)
module = CustomModule()
Explanation: The assets directory contains files used by the TensorFlow graph, for example text files used to initialize vocabulary tables. It is unused in this example.
SavedModels may have an assets.extra directory for any files not used by the TensorFlow graph, for example information for consumers about what to do with the SavedModel. TensorFlow itself does not use this directory.
Saving a custom model
tf.saved_model.save supports saving tf.Module objects and its subclasses, like tf.keras.Layer and tf.keras.Model.
Let's look at an example of saving and restoring a tf.Module.
End of explanation
module_no_signatures_path = os.path.join(tmpdir, 'module_no_signatures')
module(tf.constant(0.))
print('Saving model...')
tf.saved_model.save(module, module_no_signatures_path)
Explanation: When you save a tf.Module, any tf.Variable attributes, tf.function-decorated methods, and tf.Modules found via recursive traversal are saved. (See the Checkpoint tutorial for more about this recursive traversal.) However, any Python attributes, functions, and data are lost. This means that when a tf.function is saved, no Python code is saved.
If no Python code is saved, how does SavedModel know how to restore the function?
Briefly, tf.function works by tracing the Python code to generate a ConcreteFunction (a callable wrapper around tf.Graph). When saving a tf.function, you're really saving the tf.function's cache of ConcreteFunctions.
To learn more about the relationship between tf.function and ConcreteFunctions, refer to the tf.function guide.
End of explanation
imported = tf.saved_model.load(module_no_signatures_path)
assert imported(tf.constant(3.)).numpy() == 3
imported.mutate(tf.constant(2.))
assert imported(tf.constant(3.)).numpy() == 6
Explanation: Loading and using a custom model
When you load a SavedModel in Python, all tf.Variable attributes, tf.function-decorated methods, and tf.Modules are restored in the same object structure as the original saved tf.Module.
End of explanation
optimizer = tf.optimizers.SGD(0.05)
def train_step():
  with tf.GradientTape() as tape:
    loss = (10. - imported(tf.constant(2.))) ** 2
  variables = tape.watched_variables()
  grads = tape.gradient(loss, variables)
  optimizer.apply_gradients(zip(grads, variables))
  return loss
for _ in range(10):
  # "v" approaches 5, "loss" approaches 0
  print("loss={:.2f} v={:.2f}".format(train_step(), imported.v.numpy()))
Explanation: Because no Python code is saved, calling a tf.function with a new input signature will fail:
python
imported(tf.constant([3.]))
<pre>
ValueError: Could not find matching function to call for canonicalized inputs ((<tf.Tensor 'args_0:0' shape=(1,) dtype=float32>,), {}). Only existing signatures are [((TensorSpec(shape=(), dtype=tf.float32, name=u'x'),), {})].
</pre>
Basic fine-tuning
Variable objects are available, and you can backprop through imported functions. That is enough to fine-tune (i.e. retrain) a SavedModel in simple cases.
End of explanation
loaded = tf.saved_model.load(mobilenet_save_path)
print("MobileNet has {} trainable variables: {}, ...".format(
          len(loaded.trainable_variables),
          ", ".join([v.name for v in loaded.trainable_variables[:5]])))
trainable_variable_ids = {id(v) for v in loaded.trainable_variables}
non_trainable_variables = [v for v in loaded.variables
                           if id(v) not in trainable_variable_ids]
print("MobileNet also has {} non-trainable variables: {}, ...".format(
          len(non_trainable_variables),
          ", ".join([v.name for v in non_trainable_variables[:3]])))
Explanation: General fine-tuning
A SavedModel from Keras provides more details than a plain __call__ to address more advanced cases of fine-tuning. TensorFlow Hub recommends to provide the following of those, if applicable, in SavedModels shared for the purpose of fine-tuning:
If the model uses dropout or another technique in which the forward pass differs between training and inference (like batch normalization), the __call__ method takes an optional, Python-valued training= argument that defaults to False but can be set to True.
Next to the __call__ attribute, there are .variable and .trainable_variable attributes with the corresponding lists of variables. A variable that was originally trainable but is meant to be frozen during fine-tuning is omitted from .trainable_variables.
For the sake of frameworks like Keras that represent weight regularizers as attributes of layers or sub-models, there can also be a .regularization_losses attribute. It holds a list of zero-argument functions whose values are meant for addition to the total loss.
Going back to the initial MobileNet example, you can see some of those in action:
End of explanation
assert len(imported.signatures) == 0
Explanation: Specifying signatures during export
Tools like TensorFlow Serving and saved_model_cli can interact with SavedModels. To help these tools determine which ConcreteFunctions to use, you need to specify serving signatures. tf.keras.Models automatically specify serving signatures, but you'll have to explicitly declare a serving signature for our custom modules.
IMPORTANT: Unless you need to export your model to an environment other than TensorFlow 2.x with Python, you probably don't need to export signatures explicitly. If you're looking for a way of enforcing an input signature for a specific function, see the input_signature argument to tf.function.
By default, no signatures are declared in a custom tf.Module.
End of explanation
module_with_signature_path = os.path.join(tmpdir, 'module_with_signature')
call = module.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
tf.saved_model.save(module, module_with_signature_path, signatures=call)
imported_with_signatures = tf.saved_model.load(module_with_signature_path)
list(imported_with_signatures.signatures.keys())
Explanation: To declare a serving signature, specify a ConcreteFunction using the signatures kwarg. When specifying a single signature, its signature key will be 'serving_default', which is saved as the constant tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY.
End of explanation
module_multiple_signatures_path = os.path.join(tmpdir, 'module_with_multiple_signatures')
signatures = {"serving_default": call,
              "array_input": module.__call__.get_concrete_function(tf.TensorSpec([None], tf.float32))}
tf.saved_model.save(module, module_multiple_signatures_path, signatures=signatures)
imported_with_multiple_signatures = tf.saved_model.load(module_multiple_signatures_path)
list(imported_with_multiple_signatures.signatures.keys())
Explanation: To export multiple signatures, pass a dictionary of signature keys to ConcreteFunctions. Each signature key corresponds to one ConcreteFunction.
End of explanation
class CustomModuleWithOutputName(tf.Module):
  def __init__(self):
    super(CustomModuleWithOutputName, self).__init__()
    self.v = tf.Variable(1.)
  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def __call__(self, x):
    return {'custom_output_name': x * self.v}
module_output = CustomModuleWithOutputName()
call_output = module_output.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
module_output_path = os.path.join(tmpdir, 'module_with_output_name')
tf.saved_model.save(module_output, module_output_path,
                    signatures={'serving_default': call_output})
imported_with_output_name = tf.saved_model.load(module_output_path)
imported_with_output_name.signatures['serving_default'].structured_outputs
Explanation: By default, the output tensor names are fairly generic, like output_0. To control the names of outputs, modify your tf.function to return a dictionary that maps output names to outputs. The names of inputs are derived from the Python function arg names.
End of explanation
<END_TASK>