1, References
TensorFlow 2.0 - tf.saved_model.save model export
Use SavedModel format
2, Related introduction
1. SavedModel format
In order to deploy the trained machine learning model to each target platform (such as server, mobile terminal, embedded device and browser), our first step is often to export (serialize) the trained whole model into a series of files in standard format. On this basis, we can use corresponding deployment tools to deploy model files on different platforms.
TensorFlow provides the first mock exam SavedModel, which enables our trained models to be deployed in different platforms on this format. This is the main export format we use in TensorFlow 2.
SavedModel format contains the complete information of a TensorFlow program: it not only includes the weight of parameters, but also the calculation process (i.e. calculation diagram)** When the model is exported as a SavedModel file, the model can be run again without the source code of the model, which makes SavedModel especially suitable for model sharing and deployment** TensorFlow Serving (server-side deployment model), TensorFlow Lite (mobile side deployment model), and TensorFlow JS will use this format.
2. SavedModel format on disk
SavedModel is a directory containing serialized signatures and the States needed to run them, including variable values and vocabularies.
.
├── efficientdet-d0_frozen.pb
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index
- saved_ model. The Pb file is used to store the actual TensorFlow program or model, as well as a set of named signatures -- each signature identifies a function that accepts tensor input and generates tensor output.
- The variables directory contains a standard training checkpoint.
- The assets directory contains files used by TensorFlow calculation diagrams, such as text files used to initialize vocabularies. This file is not used in this example.
3, Save and load in SavedModel format
1. Save / export
tf.saved_model.save model export, that is, model saving
# Export model, model catalog tf.saved_model.save(mymodel, "Save to folder name")
2. Load / import
tf.saved_model.load model import, that is, model loading
# Load model
mymodel = tf.saved_model.load("Save to folder name")
4, Save custom model
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()
When saved tf.Module When any tf.Variable Characteristics tf.function Decoration method and found by recursive traversal tf.Module Will be saved. All Python features, functions, and data are lost. That is, when saving tf.function Python code is not saved when.
In short, tf.function It works by tracing Python code to generate a ConcreteFunction (a callable tf.Graph Wrapper). When saved tf.function What is actually saved is tf.function The ConcreteFunction cache for.
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)
# output
Tracing with Tensor("x:0", shape=(), dtype=float32)
Saving model...
Tracing with Tensor("x:0", shape=(), dtype=float32)
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_no_signatures/assets
5, Loading and using custom models
When SavedModel is loaded in Python, all tf.Variable Characteristics tf.function Decoration method and method tf.Module Will be saved according to the original tf.Module Restore the same object structure.
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
Since Python code is not saved, it is called with a new input signature tf.function Will fail:
# output
imported(tf.constant([3.]))
ValueError: Could not find matching function to call for canonicalized inputs ((,), {}). Only existing signatures are [((TensorSpec(shape=(), dtype=tf.float32, name=u'x'),), {})].
have access to tf.saved_model.load Load SavedModel back into Python.
loaded = tf.saved_model.load(mobilenet_save_path) print(list(loaded.signatures.keys())) # ["serving_default"] # output ["serving_default"]
Imported signatures always return to the dictionary.
infer = loaded.signatures["serving_default"]
print(infer.structured_outputs)
# output
{'predictions': TensorSpec(shape=(None, 1000), dtype=tf.float32, name='predictions')}
6, Specify signature when exporting
tf.keras.Model The service online signature will be automatically specified. However, for the custom module, we must explicitly declare the service online signature.
When a single signature is specified, the signature key is' serving '_ Default 'and save as a constant tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY.
By default, custom tf.Module Signature will not be declared in.
assert len(imported.signatures) == 0
If you specify a signature when exporting, you need to use the signatures keyword parameter to specify ConcreteFunction.
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)
# output
Tracing with Tensor("x:0", dtype=float32)
Tracing with Tensor("x:0", dtype=float32)
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_with_signature/assets
imported_with_signatures = tf.saved_model.load(module_with_signature_path) list(imported_with_signatures.signatures.keys()) # output ['serving_default']
To export multiple signatures, pass the dictionary of signature keys to ConcreteFunction. Each signature key corresponds to a ConcreteFunction.
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)
# output
Tracing with Tensor("x:0", shape=(None,), dtype=float32)
Tracing with Tensor("x:0", shape=(None,), dtype=float32)
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_with_multiple_signatures/assets
imported_with_multiple_signatures = tf.saved_model.load(module_multiple_signatures_path) list(imported_with_multiple_signatures.signatures.keys()) # output ['serving_default', 'array_input']
By default, the output tensor name is very general, such as output_0 To control the name of the output, modify the tf.function To return the dictionary that maps the output name to the output. The name entered is from the Python function parameter name.
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})
# output
INFO:tensorflow:Assets written to: /tmp/tmp0ya0f8kf/module_with_output_name/assets
imported_with_output_name = tf.saved_model.load(module_output_path)
imported_with_output_name.signatures['serving_default'].structured_outputs
# output
{'custom_output_name': TensorSpec(shape=(), dtype=tf.float32, name='custom_output_name')}
7, Keras API model export
# Model export
model.save('catdog.h5')
# Model loading
model = tf.keras.models.load_model('catdog.h5')
8, Model py
"""
because SavedModel Based on the calculation diagram, so for using inheritance tf.keras.Model Class established Keras Model, which needs to be exported to SavedModel Format method (e.g call )All need to use @tf.function modification.
"""
class MLPmodel(tf.keras.Model):
def __init__(self):
super().__init__()
# Flatten dimensions other than the first dimension
self.flatten = tf.keras.layers.Flatten()
self.dense1 = tf.keras.layers.Dense(units=100, activation='relu')
self.dense2 = tf.keras.layers.Dense(units=10)
@tf.function # Calculation diagram mode, export model, must be written
def call(self, input):
x = self.flatten(input)
x = self.dense1(x)
x = self.dense2(x)
output = tf.nn.softmax(x)
return output
# Export model, model catalog
tf.saved_model.save(mymodel, "./my_model_path")
# Load model
mymodel = tf.saved_model.load('./my_model_path')
# For using inheritance TF Keras. The Keras model created by the model class cannot be evaluate d after loading with SavedModel. Instead, it needs to use model call() . res = mymodel.call(data_loader.test_data) print(res)