wehub-resource-sync
517 lines
20 KiB
Python
517 lines
20 KiB
Python
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# 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
|
|
#
|
|
# http://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.
|
|
# ==============================================================================
|
|
"""Loader implementation for SavedModel with hermetic, language-neutral exports.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
from google.protobuf import message
|
|
from google.protobuf import text_format
|
|
|
|
from tensorflow.core.framework import graph_debug_info_pb2
|
|
from tensorflow.core.protobuf import meta_graph_pb2
|
|
from tensorflow.core.protobuf import saved_model_pb2
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.lib.io import file_io
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.platform import tf_logging
|
|
from tensorflow.python.saved_model import constants
|
|
from tensorflow.python.saved_model import path_helpers
|
|
from tensorflow.python.saved_model import signature_def_utils
|
|
from tensorflow.python.saved_model import utils_impl as saved_model_utils
|
|
# Placeholder for protosplitter merger import.
|
|
from tensorflow.python.saved_model.pywrap_saved_model import metrics
|
|
from tensorflow.python.training import saver as tf_saver
|
|
from tensorflow.python.util import compat
|
|
from tensorflow.python.util import deprecation
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
|
|
# API label for SavedModel metrics.
|
|
_LOADER_LABEL = "loader"
|
|
|
|
|
|
def parse_saved_model_with_debug_info(export_dir):
|
|
"""Reads the savedmodel as well as the graph debug info.
|
|
|
|
Args:
|
|
export_dir: Directory containing the SavedModel and GraphDebugInfo files.
|
|
|
|
Returns:
|
|
`SavedModel` and `GraphDebugInfo` protocol buffers.
|
|
|
|
Raises:
|
|
IOError: If the saved model file does not exist, or cannot be successfully
|
|
parsed. Missing graph debug info file is fine.
|
|
"""
|
|
saved_model = parse_saved_model(export_dir)
|
|
|
|
debug_info_path = file_io.join(
|
|
path_helpers.get_debug_dir(export_dir),
|
|
constants.DEBUG_INFO_FILENAME_PB)
|
|
debug_info = graph_debug_info_pb2.GraphDebugInfo()
|
|
if file_io.file_exists(debug_info_path):
|
|
with file_io.FileIO(debug_info_path, "rb") as debug_file:
|
|
try:
|
|
debug_info.ParseFromString(debug_file.read())
|
|
except message.DecodeError as e:
|
|
raise IOError(f"Cannot parse file {debug_info_path}: {e}.")
|
|
|
|
return (saved_model, debug_info)
|
|
|
|
|
|
@tf_export("__internal__.saved_model.parse_saved_model", v1=[])
|
|
def parse_saved_model(export_dir):
|
|
"""Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`.
|
|
|
|
Args:
|
|
export_dir: String or Pathlike, path to the directory containing the
|
|
SavedModel file.
|
|
|
|
Returns:
|
|
A `SavedModel` protocol buffer.
|
|
|
|
Raises:
|
|
IOError: If the file does not exist, or cannot be successfully parsed.
|
|
"""
|
|
# Build the path to the SavedModel in pbtxt format.
|
|
path_to_pbtxt = file_io.join(
|
|
compat.as_bytes(compat.path_to_str(export_dir)),
|
|
compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
|
|
# Build the path to the SavedModel in pb format.
|
|
path_to_pb = file_io.join(
|
|
compat.as_bytes(compat.path_to_str(export_dir)),
|
|
compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
|
|
# Build the path to the SavedModel in cpb format.
|
|
path_to_cpb = file_io.join(
|
|
compat.as_bytes(compat.path_to_str(export_dir)),
|
|
compat.as_bytes(constants.SAVED_MODEL_FILENAME_CPB))
|
|
|
|
# Parse the SavedModel protocol buffer.
|
|
saved_model = saved_model_pb2.SavedModel()
|
|
if file_io.file_exists(path_to_pb):
|
|
with file_io.FileIO(path_to_pb, "rb") as f:
|
|
file_content = f.read()
|
|
try:
|
|
saved_model.ParseFromString(file_content)
|
|
except message.DecodeError as e:
|
|
raise IOError(f"Cannot parse file {path_to_pb}: {str(e)}.") from e
|
|
elif file_io.file_exists(path_to_pbtxt):
|
|
with file_io.FileIO(path_to_pbtxt, "rb") as f:
|
|
file_content = f.read()
|
|
try:
|
|
text_format.Parse(file_content.decode("utf-8"), saved_model)
|
|
except text_format.ParseError as e:
|
|
raise IOError(f"Cannot parse file {path_to_pbtxt}: {str(e)}.") from e
|
|
else:
|
|
raise IOError(
|
|
f"SavedModel file does not exist at: {export_dir}{os.path.sep}"
|
|
f"{{{constants.SAVED_MODEL_FILENAME_PBTXT}|"
|
|
f"{constants.SAVED_MODEL_FILENAME_PB}}}")
|
|
return saved_model
|
|
|
|
|
|
def get_asset_tensors(export_dir, meta_graph_def_to_load, import_scope=None):
|
|
"""Gets the asset tensors, if defined in the meta graph def to load.
|
|
|
|
Args:
|
|
export_dir: Directory where the SavedModel is located.
|
|
meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.
|
|
import_scope: Optional `string` -- if specified, prepend this followed by
|
|
'/' to all returned asset tensor names.
|
|
|
|
Returns:
|
|
A dictionary of asset tensors, keyed by the name of the asset tensor. The
|
|
value in the map corresponds to the absolute path of the asset file.
|
|
"""
|
|
# Collection-def that may contain the assets key.
|
|
collection_def = meta_graph_def_to_load.collection_def
|
|
|
|
asset_tensor_dict = {}
|
|
asset_protos = []
|
|
|
|
if meta_graph_def_to_load.asset_file_def:
|
|
asset_protos = meta_graph_def_to_load.asset_file_def
|
|
elif constants.ASSETS_KEY in collection_def:
|
|
assets_any_proto = collection_def[constants.ASSETS_KEY].any_list.value
|
|
for asset_any_proto in assets_any_proto:
|
|
asset_proto = meta_graph_pb2.AssetFileDef()
|
|
asset_any_proto.Unpack(asset_proto)
|
|
asset_protos.append(asset_proto)
|
|
|
|
# Location of the assets for SavedModel.
|
|
assets_directory = file_io.join(
|
|
compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY))
|
|
# Process each asset and add it to the asset tensor dictionary.
|
|
for asset_proto in asset_protos:
|
|
tensor_name = asset_proto.tensor_info.name
|
|
if import_scope:
|
|
tensor_name = "%s/%s" % (import_scope, tensor_name)
|
|
asset_tensor_dict[tensor_name] = file_io.join(
|
|
compat.as_bytes(assets_directory),
|
|
compat.as_bytes(asset_proto.filename))
|
|
|
|
return asset_tensor_dict
|
|
|
|
|
|
def _get_main_op_tensor(
|
|
meta_graph_def_to_load, init_op_key=constants.MAIN_OP_KEY):
|
|
"""Gets the main op tensor, if one exists.
|
|
|
|
Args:
|
|
meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.
|
|
init_op_key: name of the collection to check; should be one of MAIN_OP_KEY
|
|
or the deprecated LEGACY_INIT_OP_KEY
|
|
|
|
Returns:
|
|
The main op tensor, if it exists and `None` otherwise.
|
|
|
|
Raises:
|
|
RuntimeError: If the collection def corresponding to the main op key has
|
|
other than exactly one tensor.
|
|
"""
|
|
collection_def = meta_graph_def_to_load.collection_def
|
|
init_op = None
|
|
if init_op_key in collection_def:
|
|
init_op_list = collection_def[init_op_key].node_list.value
|
|
if len(init_op_list) != 1:
|
|
raise RuntimeError("Expected exactly one SavedModel init op. "
|
|
f"Found {len(init_op_list)}: {init_op_list}.")
|
|
init_op = ops.get_collection(init_op_key)[0]
|
|
return init_op
|
|
|
|
|
|
def _get_op_from_collection(meta_graph_def, op_key):
|
|
return _get_main_op_tensor(meta_graph_def, op_key)
|
|
|
|
|
|
def _get_op_from_signature_def(meta_graph_def, op_signature_key, import_scope):
|
|
"""Retrieve op stored in the imported meta graph's signature def."""
|
|
if op_signature_key in meta_graph_def.signature_def:
|
|
return signature_def_utils.load_op_from_signature_def(
|
|
meta_graph_def.signature_def[op_signature_key], op_signature_key,
|
|
import_scope)
|
|
else:
|
|
return None
|
|
|
|
|
|
def get_init_op(meta_graph_def, import_scope=None):
|
|
return (_get_op_from_signature_def(
|
|
meta_graph_def, constants.INIT_OP_SIGNATURE_KEY, import_scope) or
|
|
_get_op_from_collection(meta_graph_def, constants.MAIN_OP_KEY) or
|
|
_get_op_from_collection(meta_graph_def, constants.LEGACY_INIT_OP_KEY))
|
|
|
|
|
|
def get_train_op(meta_graph_def, import_scope=None):
|
|
train_op = _get_op_from_signature_def(
|
|
meta_graph_def, constants.TRAIN_OP_SIGNATURE_KEY, import_scope)
|
|
if train_op is None:
|
|
train_op = _get_op_from_collection(meta_graph_def, constants.TRAIN_OP_KEY)
|
|
return train_op
|
|
|
|
|
|
@tf_export(v1=[
|
|
"saved_model.contains_saved_model",
|
|
"saved_model.maybe_saved_model_directory",
|
|
"saved_model.loader.maybe_saved_model_directory"
|
|
])
|
|
@deprecation.deprecated_endpoints(
|
|
"saved_model.loader.maybe_saved_model_directory")
|
|
def maybe_saved_model_directory(export_dir):
|
|
"""Checks whether the provided export directory could contain a SavedModel.
|
|
|
|
Note that the method does not load any data by itself. If the method returns
|
|
`false`, the export directory definitely does not contain a SavedModel. If the
|
|
method returns `true`, the export directory may contain a SavedModel but
|
|
provides no guarantee that it can be loaded.
|
|
|
|
Args:
|
|
export_dir: Absolute string path to possible export location. For example,
|
|
'/my/foo/model'.
|
|
|
|
Returns:
|
|
True if the export directory contains SavedModel files, False otherwise.
|
|
"""
|
|
txt_path = file_io.join(export_dir, constants.SAVED_MODEL_FILENAME_PBTXT)
|
|
pb_path = file_io.join(export_dir, constants.SAVED_MODEL_FILENAME_PB)
|
|
cpb_path = file_io.join(export_dir, constants.SAVED_MODEL_FILENAME_CPB)
|
|
return (
|
|
file_io.file_exists(txt_path)
|
|
or file_io.file_exists(pb_path)
|
|
or file_io.file_exists(cpb_path)
|
|
)
|
|
|
|
|
|
@tf_export("saved_model.contains_saved_model", v1=[])
|
|
def contains_saved_model(export_dir):
|
|
"""Checks whether the provided export directory could contain a SavedModel.
|
|
|
|
Note that the method does not load any data by itself. If the method returns
|
|
`false`, the export directory definitely does not contain a SavedModel. If the
|
|
method returns `true`, the export directory may contain a SavedModel but
|
|
provides no guarantee that it can be loaded.
|
|
|
|
Args:
|
|
export_dir: Absolute path to possible export location. For example,
|
|
'/my/foo/model'.
|
|
|
|
Returns:
|
|
True if the export directory contains SavedModel files, False otherwise.
|
|
"""
|
|
if isinstance(export_dir, os.PathLike):
|
|
export_dir = os.fspath(export_dir)
|
|
return maybe_saved_model_directory(export_dir)
|
|
|
|
|
|
@tf_export(v1=["saved_model.load", "saved_model.loader.load"])
|
|
@deprecation.deprecated(
|
|
None,
|
|
"Use `tf.saved_model.load` instead.")
|
|
def load(sess, tags, export_dir, import_scope=None, **saver_kwargs):
|
|
"""Loads the model from a SavedModel as specified by tags.
|
|
|
|
Args:
|
|
sess: The TensorFlow session to restore the variables.
|
|
tags: Set of string tags to identify the required MetaGraphDef. These should
|
|
correspond to the tags used when saving the variables using the
|
|
SavedModel `save()` API.
|
|
export_dir: Directory in which the SavedModel protocol buffer and variables
|
|
to be loaded are located.
|
|
import_scope: Optional `string` -- if specified, prepend this string
|
|
followed by '/' to all loaded tensor names. This scope is applied to
|
|
tensor instances loaded into the passed session, but it is *not* written
|
|
through to the static `MetaGraphDef` protocol buffer that is returned.
|
|
**saver_kwargs: Optional keyword arguments passed through to Saver.
|
|
|
|
Returns:
|
|
The `MetaGraphDef` protocol buffer loaded in the provided session. This
|
|
can be used to further extract signature-defs, collection-defs, etc.
|
|
|
|
Raises:
|
|
RuntimeError: MetaGraphDef associated with the tags cannot be found.
|
|
|
|
@compatibility(TF2)
|
|
|
|
`tf.compat.v1.saved_model.load` or `tf.compat.v1.saved_model.loader.load` is
|
|
not compatible with eager execution. Please use `tf.saved_model.load` instead
|
|
to load your model. You can refer to the [SavedModel guide]
|
|
(https://www.tensorflow.org/guide/saved_model) for more information as well as
|
|
"Importing SavedModels from TensorFlow 1.x" in the [`tf.saved_model.load`]
|
|
(https://www.tensorflow.org/api_docs/python/tf/saved_model/load) docstring.
|
|
|
|
#### How to Map Arguments
|
|
|
|
| TF1 Arg Name | TF2 Arg Name | Note |
|
|
| :-------------------- | :-------------- | :------------------------- |
|
|
| `sess` | Not supported | - |
|
|
| `tags` | `tags` | - |
|
|
| `export_dir` | `export_dir` | - |
|
|
| `import_scope` | Not supported | Name scopes are not needed.
|
|
: : : By default, variables are :
|
|
: : : associated with the loaded :
|
|
: : : object and function names :
|
|
: : : are deduped. :
|
|
| `saver_kwargs` | Not supported | - |
|
|
|
|
#### Before & After Usage Example
|
|
|
|
Before:
|
|
|
|
```
|
|
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
|
|
tf.compat.v1.saved_model.loader.load(sess, ["foo-tag"], export_dir)
|
|
```
|
|
|
|
After:
|
|
|
|
```
|
|
model = tf.saved_model.load(export_dir, tags=["foo-tag"])
|
|
```
|
|
@end_compatibility
|
|
"""
|
|
loader = SavedModelLoader(export_dir)
|
|
return loader.load(sess, tags, import_scope, **saver_kwargs)
|
|
|
|
|
|
class SavedModelLoader(object):
|
|
"""Load graphs and restore variable values from a `SavedModel`."""
|
|
|
|
def __init__(self, export_dir):
|
|
"""Creates a `SavedModelLoader`.
|
|
|
|
Args:
|
|
export_dir: Directory in which the SavedModel protocol buffer and
|
|
variables to be loaded are located.
|
|
"""
|
|
self._export_dir = export_dir
|
|
self._variables_path = path_helpers.get_variables_path(export_dir)
|
|
self._saved_model = parse_saved_model(export_dir)
|
|
|
|
@property
|
|
def export_dir(self):
|
|
"""Directory containing the SavedModel."""
|
|
return self._export_dir
|
|
|
|
@property
|
|
def variables_path(self):
|
|
"""Path to variable checkpoint files."""
|
|
return self._variables_path
|
|
|
|
@property
|
|
def saved_model(self):
|
|
"""SavedModel object parsed from the export directory."""
|
|
return self._saved_model
|
|
|
|
def get_meta_graph_def_from_tags(self, tags):
|
|
"""Return MetaGraphDef with the exact specified tags.
|
|
|
|
Args:
|
|
tags: A list or set of string tags that identify the MetaGraphDef.
|
|
|
|
Returns:
|
|
MetaGraphDef with the same tags.
|
|
|
|
Raises:
|
|
RuntimeError: if no metagraphs were found with the associated tags.
|
|
"""
|
|
found_match = False
|
|
meta_graph_def_to_load = None
|
|
available_tags = []
|
|
for meta_graph_def in self._saved_model.meta_graphs:
|
|
available_tags.append(set(meta_graph_def.meta_info_def.tags))
|
|
if set(meta_graph_def.meta_info_def.tags) == set(tags):
|
|
meta_graph_def_to_load = meta_graph_def
|
|
found_match = True
|
|
break
|
|
|
|
if not found_match:
|
|
raise RuntimeError(
|
|
f"MetaGraphDef associated with tags {str(tags).strip('[]')} "
|
|
"could not be found in SavedModel, with available tags "
|
|
f"'{available_tags}'. To inspect available tag-sets in"
|
|
" the SavedModel, please use the SavedModel CLI: `saved_model_cli`.")
|
|
return meta_graph_def_to_load
|
|
|
|
def load_graph(self, graph, tags, import_scope=None, **saver_kwargs):
|
|
"""Load ops and nodes from SavedModel MetaGraph into graph.
|
|
|
|
Args:
|
|
graph: tf.Graph object.
|
|
tags: a set of string tags identifying a MetaGraphDef.
|
|
import_scope: Optional `string` -- if specified, prepend this string
|
|
followed by '/' to all loaded tensor names. This scope is applied to
|
|
tensor instances loaded into the passed session, but it is *not* written
|
|
through to the static `MetaGraphDef` protocol buffer that is returned.
|
|
**saver_kwargs: keyword arguments to pass to tf.train.import_meta_graph.
|
|
|
|
Returns:
|
|
A tuple of
|
|
* Saver defined by the MetaGraph, which can be used to restore the
|
|
variable values.
|
|
* List of `Operation`/`Tensor` objects returned from
|
|
`tf.import_graph_def` (may be `None`).
|
|
"""
|
|
meta_graph_def = self.get_meta_graph_def_from_tags(tags)
|
|
if sys.byteorder == "big":
|
|
saved_model_utils.swap_function_tensor_content(meta_graph_def, "little",
|
|
"big")
|
|
with graph.as_default():
|
|
return tf_saver._import_meta_graph_with_return_elements( # pylint: disable=protected-access
|
|
meta_graph_def, import_scope=import_scope, **saver_kwargs)
|
|
|
|
def restore_variables(self, sess, saver, import_scope=None):
|
|
"""Restore SavedModel variable values into the session.
|
|
|
|
Args:
|
|
sess: tf.compat.v1.Session to restore variable values.
|
|
saver: a tf.compat.v1.train.Saver object. Can be None if there are no
|
|
variables in graph. This may be the saver returned by the load_graph()
|
|
function, or a default `tf.compat.v1.train.Saver()`.
|
|
import_scope: Optional `string` -- if specified, prepend this string
|
|
followed by '/' to all loaded tensor names. This scope is applied to
|
|
tensor instances loaded into the passed session, but it is *not* written
|
|
through to the static `MetaGraphDef` protocol buffer that is returned.
|
|
|
|
Raises:
|
|
ValueError: if no saver was passed to the saver argument, and there are
|
|
variables in the graph.
|
|
"""
|
|
with sess.graph.as_default():
|
|
if (saver is None and
|
|
not variables._all_saveable_objects(scope=import_scope)): # pylint: disable=protected-access
|
|
tf_logging.info("The specified SavedModel has no variables; no "
|
|
"checkpoints were restored.")
|
|
elif isinstance(saver, tf_saver.Saver):
|
|
saver.restore(sess, self._variables_path)
|
|
else:
|
|
raise ValueError(
|
|
"No tf.train.Saver object was passed to the function "
|
|
"`SavedModelLoader.restore_variables`. Since there are variables in"
|
|
" the graph, a saver is required.")
|
|
|
|
def run_init_ops(self, sess, tags, import_scope=None):
|
|
"""Run initialization ops defined in the `MetaGraphDef`.
|
|
|
|
Args:
|
|
sess: tf.compat.v1.Session to restore variable values.
|
|
tags: a set of string tags identifying a MetaGraphDef.
|
|
import_scope: Optional `string` -- if specified, prepend this string
|
|
followed by '/' to all loaded tensor names. This scope is applied to
|
|
tensor instances loaded into the passed session, but it is *not* written
|
|
through to the static `MetaGraphDef` protocol buffer that is returned.
|
|
"""
|
|
meta_graph_def = self.get_meta_graph_def_from_tags(tags)
|
|
with sess.graph.as_default():
|
|
# Get asset tensors, if any.
|
|
asset_tensors_dictionary = get_asset_tensors(
|
|
self._export_dir, meta_graph_def, import_scope=import_scope)
|
|
|
|
init_op = get_init_op(meta_graph_def, import_scope)
|
|
if init_op is not None:
|
|
sess.run(fetches=[init_op], feed_dict=asset_tensors_dictionary)
|
|
|
|
def load(self, sess, tags, import_scope=None, **saver_kwargs):
|
|
"""Load the MetaGraphDef graph and restore variable values into the session.
|
|
|
|
Args:
|
|
sess: tf.compat.v1.Session to restore variable values.
|
|
tags: a set of string tags identifying a MetaGraphDef.
|
|
import_scope: Optional `string` -- if specified, prepend this string
|
|
followed by '/' to all loaded tensor names. This scope is applied to
|
|
tensor instances loaded into the passed session, but it is *not* written
|
|
through to the static `MetaGraphDef` protocol buffer that is returned.
|
|
**saver_kwargs: keyword arguments to pass to tf.train.import_meta_graph.
|
|
|
|
Returns:
|
|
`MetagraphDef` proto of the graph that was loaded.
|
|
"""
|
|
saved_model_proto = parse_saved_model(self._export_dir)
|
|
metrics.IncrementReadApi(_LOADER_LABEL)
|
|
|
|
with sess.graph.as_default():
|
|
saver, _ = self.load_graph(sess.graph, tags, import_scope,
|
|
**saver_kwargs)
|
|
self.restore_variables(sess, saver, import_scope)
|
|
self.run_init_ops(sess, tags, import_scope)
|
|
meta_graph_def = self.get_meta_graph_def_from_tags(tags)
|
|
|
|
if (len(saved_model_proto.meta_graphs) == 1 and
|
|
saved_model_proto.meta_graphs[0].HasField("object_graph_def")):
|
|
metrics.IncrementRead(write_version="2")
|
|
else:
|
|
metrics.IncrementRead(write_version="1")
|
|
|
|
return meta_graph_def
|