View source on GitHub |
ExportArchive is used to write SavedModel artifacts (e.g. for inference).
tf.keras.export.ExportArchive()
If you have a Keras model or layer that you want to export as SavedModel for
serving (e.g. via TensorFlow-Serving), you can use ExportArchive
to configure the different serving endpoints you need to make available,
as well as their signatures. Simply instantiate an ExportArchive
,
use track()
to register the layer(s) or model(s) to be used,
then use the add_endpoint()
method to register a new serving endpoint.
When done, use the write_out()
method to save the artifact.
The resulting artifact is a SavedModel and can be reloaded via
tf.saved_model.load
.
Examples:
Here's how to export a model for inference.
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
name="serve",
fn=model.call,
input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
)
export_archive.write_out("path/to/location")
# Elsewhere, we can reload the artifact and serve it.
# The endpoint we added is available as a method:
serving_model = tf.saved_model.load("path/to/location")
outputs = serving_model.serve(inputs)
Here's how to export a model with one endpoint for inference and one endpoint for a training-mode forward pass (e.g. with dropout on).
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
name="call_inference",
fn=lambda x: model.call(x, training=False),
input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
)
export_archive.add_endpoint(
name="call_training",
fn=lambda x: model.call(x, training=True),
input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
)
export_archive.write_out("path/to/location")
Methods
add_endpoint
add_endpoint(
name, fn, input_signature=None
)
Register a new serving endpoint.
Arguments | |
---|---|
name
|
Str, name of the endpoint. |
fn
|
A function. It should only leverage resources
(e.g. tf.Variable objects or tf.lookup.StaticHashTable
objects) that are available on the models/layers
tracked by the ExportArchive (you can call .track(model)
to track a new model).
The shape and dtype of the inputs to the function must be
known. For that purpose, you can either 1) make sure that
fn is a tf.function that has been called at least once, or
2) provide an |
input_signature
|
Used to specify the shape and dtype of the
inputs to fn . List of tf.TensorSpec objects (one
per positional input argument of fn ). Nested arguments are
allowed (see below for an example showing a Functional model
with 2 input arguments).
|
Example:
Adding an endpoint using the input_signature
argument when the
model has a single input argument:
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
name="serve",
fn=model.call,
input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
)
Adding an endpoint using the input_signature
argument when the
model has two positional input arguments:
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
name="serve",
fn=model.call,
input_signature=[
tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
],
)
Adding an endpoint using the input_signature
argument when the
model has one input argument that is a list of 2 tensors (e.g.
a Functional model with 2 inputs):
model = keras.Model(inputs=[x1, x2], outputs=outputs)
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
name="serve",
fn=model.call,
input_signature=[
[
tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
],
],
)
This also works with dictionary inputs:
model = keras.Model(inputs={"x1": x1, "x2": x2}, outputs=outputs)
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(
name="serve",
fn=model.call,
input_signature=[
{
"x1": tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
"x2": tf.TensorSpec(shape=(None, 4), dtype=tf.float32),
},
],
)
Adding an endpoint that is a tf.function
:
@tf.function()
def serving_fn(x):
return model(x)
# The function must be traced, i.e. it must be called at least once.
serving_fn(tf.random.normal(shape=(2, 3)))
export_archive = ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(name="serve", fn=serving_fn)
add_variable_collection
add_variable_collection(
name, variables
)
Register a set of variables to be retrieved after reloading.
Arguments | |
---|---|
name
|
The string name for the collection. |
variables
|
A tuple/list/set of tf.Variable instances.
|
Example:
export_archive = ExportArchive()
export_archive.track(model)
# Register an endpoint
export_archive.add_endpoint(
name="serve",
fn=model.call,
input_signature=[tf.TensorSpec(shape=(None, 3), dtype=tf.float32)],
)
# Save a variable collection
export_archive.add_variable_collection(
name="optimizer_variables", variables=model.optimizer.variables)
export_archive.write_out("path/to/location")
# Reload the object
revived_object = tf.saved_model.load("path/to/location")
# Retrieve the variables
optimizer_variables = revived_object.optimizer_variables
track
track(
layer
)
Track the variables (and other resources) of a layer or model.
write_out
write_out(
filepath, options=None
)
Write the corresponding SavedModel to disk.
Arguments | |
---|---|
filepath
|
str or pathlib.Path object.
Path where to save the artifact.
|
options
|
tf.saved_model.SaveOptions object that specifies
SavedModel saving options.
|
Note on TF-Serving: all endpoints registered via add_endpoint()
are made visible for TF-Serving in the SavedModel artifact. In addition,
the first endpoint registered is made visible under the alias
"serving_default"
(unless an endpoint with the name
"serving_default"
was already registered manually),
since TF-Serving requires this endpoint to be set.