def load_model(sess, path, mode):
tags = model_utils.EXPORT_TAG_MAP[mode]
sig_def_key = model_utils.SIGNATURE_KEY_MAP[mode]
meta_graph_def = loader_impl.load(sess, tags, path)
inputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()}
outputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()}
return inputs, outputs, meta_graph_def
def load_model(sess, path, mode):
tags = model_fn_lib.EXPORT_TAG_MAP[mode]
sig_def_key = (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
if mode == model_fn_lib.ModeKeys.PREDICT else mode)
meta_graph_def = loader_impl.load(sess, tags, path)
inputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()}
outputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()}
return inputs, outputs, meta_graph_def
def load_model(sess, path, mode):
tags = model_utils.EXPORT_TAG_MAP[mode]
sig_def_key = model_utils.SIGNATURE_KEY_MAP[mode]
meta_graph_def = loader_impl.load(sess, tags, path)
inputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()}
outputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()}
return inputs, outputs, meta_graph_def
def load_model(sess, path, mode):
tags = model_fn_lib.EXPORT_TAG_MAP[mode]
sig_def_key = (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
if mode == model_fn_lib.ModeKeys.PREDICT else mode)
meta_graph_def = loader_impl.load(sess, tags, path)
inputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].inputs.items()}
outputs = {
k: sess.graph.get_tensor_by_name(v.name)
for k, v in meta_graph_def.signature_def[sig_def_key].outputs.items()}
return inputs, outputs
def _run_graph_for_calibration_graph_mode(
model_dir: str, signature_keys: List[str], tags: Set[str],
representative_dataset: repr_dataset.RepresentativeDataset) -> None:
"""Runs the graph for calibration in graph mode.
This function assumes _graph mode_ (used when legacy TF1 is used or when eager
mode is explicitly disabled) when running the graph. This step is used in
order to collect the statistics in CustomAggregatorOp for quantization using
the representative dataset for the actual data provided for inference.
Args:
model_dir: Path to SavedModel directory.
signature_keys: A list of signature keys that identifies a function to run
the data samples with.
tags: Set of tags identifying the MetaGraphDef within the SavedModel.
representative_dataset: Representative dataset used for calibration.
Raises:
ValueError: When the samples in representative dataset is invalid.
"""
with session.Session() as sess:
meta_graph: meta_graph_pb2.MetaGraphDef = saved_model_loader.load(
sess, tags, export_dir=model_dir)
for sample in representative_dataset:
signature_key, input_data = _get_signature_key_and_input(
sample, signature_keys)
sig_def = meta_graph.signature_def[signature_key]
output_tensor_names = [
output_tensor_info.name
for output_tensor_info in sig_def.outputs.values()
]
# Create a mapping from input tensor name to the input tensor value.
# ex) "Placeholder:0" -> [0, 1, 2]
try:
feed_dict = _create_feed_dict_from_input_data(
input_data, sig_def)
except KeyError as key_error:
raise ValueError(
f'Invalid input data for signature: {signature_key}.'
) from key_error
sess.run(output_tensor_names, feed_dict=feed_dict)
def _run_graph_for_calibration_graph_mode(
model_dir: str,
tags: Collection[str],
representative_dataset_map: repr_dataset.RepresentativeDatasetMapping,
) -> None:
"""Runs the graph for calibration in graph mode.
This function assumes _graph mode_ (used when legacy TF1 is used or when eager
mode is explicitly disabled) when running the graph. This step is used in
order to collect the statistics in CustomAggregatorOp for quantization using
the representative dataset for the actual data provided for inference.
Args:
model_dir: Path to SavedModel directory.
tags: Collection of tags identifying the MetaGraphDef within the SavedModel.
representative_dataset_map: A map where signature keys are mapped to
corresponding representative datasets.
Raises:
ValueError: When running the function with the representative dataset fails.
"""
# Replace tf.Tensors by numpy ndarrays in order to reuse the samples in a
# different graph when running the calibration.
_replace_tensors_by_numpy_ndarrays(representative_dataset_map)
# Run the calibration in a new graph to avoid name collision, which could
# happen when the same model is loaded multiple times in the default graph.
with ops.Graph().as_default(), session.Session() as sess:
meta_graph: meta_graph_pb2.MetaGraphDef = saved_model_loader.load(
sess, tags, export_dir=model_dir)
for signature_key, repr_ds in representative_dataset_map.items():
sig_def = meta_graph.signature_def[signature_key]
try:
_run_function_for_calibration_graph_mode(
sess, signature_def=sig_def, representative_dataset=repr_ds)
except Exception as ex:
raise ValueError(
'Failed to run representative dataset through the '
f'function with the signature key: {signature_key}.') from ex
def _run_graph_for_calibration_graph_mode(
model_dir: str,
tags: Collection[str],
representative_dataset_map: repr_dataset.RepresentativeDatasetMapping,
) -> None:
"""Runs the graph for calibration in graph mode.
This function assumes _graph mode_ (used when legacy TF1 is used or when eager
mode is explicitly disabled) when running the graph. This step is used in
order to collect the statistics in CustomAggregatorOp for quantization using
the representative dataset for the actual data provided for inference.
Args:
model_dir: Path to SavedModel directory.
tags: Collection of tags identifying the MetaGraphDef within the SavedModel.
representative_dataset_map: A map where signature keys are mapped to
corresponding representative datasets.
Raises:
ValueError: When running the function with the representative dataset fails.
"""
with session.Session() as sess:
meta_graph: meta_graph_pb2.MetaGraphDef = saved_model_loader.load(
sess, tags, export_dir=model_dir)
for signature_key, repr_ds in representative_dataset_map.items():
sig_def = meta_graph.signature_def[signature_key]
try:
_run_function_for_calibration_graph_mode(
sess,
signature_def=sig_def,
representative_dataset=repr_ds)
except Exception as ex:
raise ValueError(
'Failed to run representative dataset through the '
f'function with the signature key: {signature_key}.'
) from ex
