def serving_input_receiver_fn():
"""Create the ServingInputReceiver to export a saved model.
Returns:
An instance of ServingInputReceiver.
"""
# We only assume one input, a string which containes the serialized proto.
receiver_tensors = {
'input_example_tensor':
tf.placeholder(dtype=tf.string,
shape=[None],
name='input_example_tensor')
}
# We have to filter our specs since only required tensors are
# used for inference time.
flat_feature_spec = tensorspec_utils.flatten_spec_structure(
self._feature_spec)
# We need to freeze the conditioning and inference shapes.
for key, value in flat_feature_spec.condition.items():
ref_shape = value.shape.as_list()
shape = [self._num_condition_samples_per_task] + ref_shape[1:]
flat_feature_spec.condition[key] = (
tensorspec_utils.ExtendedTensorSpec.from_spec(value,
shape=shape))
for key, value in flat_feature_spec.inference.items():
ref_shape = value.shape.as_list()
shape = [self._num_inference_samples_per_task] + ref_shape[1:]
flat_feature_spec.inference[key] = (
tensorspec_utils.ExtendedTensorSpec.from_spec(value,
shape=shape))
required_feature_spec = (
tensorspec_utils.filter_required_flat_tensor_spec(
flat_feature_spec))
tensor_dict, tensor_spec_dict = (
tensorspec_utils.tensorspec_to_feature_dict(
required_feature_spec))
parse_tf_example_fn = tfdata.create_parse_tf_example_fn(
tensor_dict=tensor_dict,
tensor_spec_dict=tensor_spec_dict,
feature_tspec=self._feature_spec)
features = parse_tf_example_fn(
receiver_tensors['input_example_tensor'])
if self._preprocess_fn is not None:
features, _ = self._preprocess_fn(
features=features,
labels=None,
mode=tf.estimator.ModeKeys.PREDICT)
return tf.estimator.export.ServingInputReceiver(
features, receiver_tensors)
def _generate_assets(model, export_dir):
in_feature_spec = model.get_feature_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
in_label_spec = model.get_label_specification_for_packing(
mode=tf.compat.v1.estimator.ModeKeys.PREDICT)
in_feature_spec = tensorspec_utils.filter_required_flat_tensor_spec(
in_feature_spec)
in_label_spec = tensorspec_utils.filter_required_flat_tensor_spec(
in_label_spec)
t2r_assets = t2r_pb2.T2RAssets()
t2r_assets.feature_spec.CopyFrom(in_feature_spec.to_proto())
t2r_assets.label_spec.CopyFrom(in_label_spec.to_proto())
t2r_assets_dir = os.path.join(export_dir,
tensorspec_utils.EXTRA_ASSETS_DIRECTORY)
tf.io.gfile.makedirs(t2r_assets_dir)
t2r_assets_filename = os.path.join(t2r_assets_dir,
tensorspec_utils.T2R_ASSETS_FILENAME)
tensorspec_utils.write_t2r_assets_to_file(t2r_assets, t2r_assets_filename)
def get_out_label_specification(self, mode):
"""The specification for the output labels after executing preprocess_fn.
Note, we strip all optional specs to further reduce communication and
computation overhead for feeding to TPUs.
Arguments:
mode: mode key for this feature specification
Returns:
A TensorSpecStruct describing the required and optional tensors.
"""
return tensorspec_utils.replace_dtype(
tensorspec_utils.filter_required_flat_tensor_spec(
self._preprocessor.get_out_label_specification(mode)),
from_dtype=tf.float32,
to_dtype=tf.bfloat16)
def serving_input_receiver_fn():
"""Create the ServingInputReceiver to export a saved model.
Returns:
An instance of ServingInputReceiver.
"""
# We have to filter our specs since only required tensors are
# used for inference time.
flat_feature_spec = tensorspec_utils.flatten_spec_structure(
self._feature_spec)
# We need to freeze the conditioning and inference shapes.
for key, value in flat_feature_spec.condition.items():
ref_shape = value.shape.as_list()
shape = [self._num_condition_samples_per_task] + ref_shape[1:]
flat_feature_spec.condition[key] = (
tensorspec_utils.ExtendedTensorSpec.from_spec(value,
shape=shape))
for key, value in flat_feature_spec.inference.items():
ref_shape = value.shape.as_list()
shape = [self._num_inference_samples_per_task] + ref_shape[1:]
flat_feature_spec.inference[key] = (
tensorspec_utils.ExtendedTensorSpec.from_spec(value,
shape=shape))
required_feature_spec = (
tensorspec_utils.filter_required_flat_tensor_spec(
flat_feature_spec))
receiver_tensors = tensorspec_utils.make_placeholders(
required_feature_spec)
# We want to ensure that our feature processing pipeline operates on a
# copy of the features and does not alter the receiver_tensors.
features = tensorspec_utils.flatten_spec_structure(
copy.copy(receiver_tensors))
if self._preprocess_fn is not None:
features, _ = self._preprocess_fn(
features=features,
labels=None,
mode=tf.estimator.ModeKeys.PREDICT)
return tf.estimator.export.ServingInputReceiver(
features, receiver_tensors)
def serving_input_receiver_fn():
"""Create the ServingInputReceiver to export a saved model.
Returns:
An instance of ServingInputReceiver.
"""
# We assume one input (a string which containes the serialized proto) per
# dataset_key.
feature_spec = self._get_input_features_for_receiver_fn()
# We have to filter our specs since only required tensors are
# used for inference time.
flat_feature_spec = tensorspec_utils.flatten_spec_structure(
feature_spec)
required_feature_spec = (
tensorspec_utils.filter_required_flat_tensor_spec(
flat_feature_spec))
dataset_keys = set(
[t.dataset_key for t in required_feature_spec.values()])
receiver_tensors = {}
parse_tensors = {}
for dataset_key in dataset_keys:
receiver_name = 'input_example_' + six.ensure_str(
(dataset_key or 'tensor'))
parse_tensors[dataset_key] = tf.placeholder(dtype=tf.string,
shape=[None],
name=receiver_name)
receiver_tensors[receiver_name] = parse_tensors[dataset_key]
parse_tf_example_fn = tfdata.create_parse_tf_example_fn(
feature_tspec=required_feature_spec)
features = parse_tf_example_fn(parse_tensors)
if (not self._export_raw_receivers
and self._preprocess_fn is not None):
features, _ = self._preprocess_fn(features=features,
labels=None)
return tf.estimator.export.ServingInputReceiver(
features, receiver_tensors)
def __init__(self,
t2r_model,
checkpoint_dir=None,
use_gpu=True,
timeout=600,
tf_intra_op_parallelism_threads=4,
tf_inter_op_parallelism_threads=4):
"""Load the model from registry and build the model in a new tf graph.
Args:
t2r_model: A T2RModel instance.
checkpoint_dir: The directory to find the checkpoint. If set to `None`, no
checkpoint will be loaded and if init_with_random_variables is set to
True a random model is initialized. Note, either checkpoint_dir or
init_with_random_variable has to be set but not both.
use_gpu: If True, will attempt to use GPU for inference.
timeout: (defaults to 600 seconds) If no checkpoint has been found after
timeout seconds restore fails.
tf_intra_op_parallelism_threads: see `tf.ConfigProto`
tf_inter_op_parallelism_threads: see `tf.ConfigProto`
"""
self._checkpoint_dir = checkpoint_dir
self._timeout = timeout
# As done in model_inference.py, a separate graph is used to build the
# target network.
g = tf.Graph()
mode = tf.estimator.ModeKeys.PREDICT
with g.as_default():
preprocessor = t2r_model.preprocessor
feature_tspec = preprocessor.get_in_feature_specification(mode)
# We perform inference, hence we only want the required tensors.
self._feature_tspec = tensorspec_utils.filter_required_flat_tensor_spec(
feature_tspec)
label_tspec = preprocessor.get_in_feature_specification(mode)
self._label_tspec = tensorspec_utils.filter_required_flat_tensor_spec(
label_tspec)
self._features = tensorspec_utils.make_placeholders(
self._feature_tspec, batch_size=None)
preprocessed_features, _ = preprocessor.preprocess(
features=self._features, labels=None, mode=mode)
estimator_spec = t2r_model.model_fn(preprocessed_features, None,
mode)
self._predictions = estimator_spec.predictions
config = tf.ConfigProto(
device_count={'GPU': 1 if use_gpu else 0},
intra_op_parallelism_threads=tf_intra_op_parallelism_threads,
inter_op_parallelism_threads=tf_inter_op_parallelism_threads)
self._sess = tf.Session(graph=g, config=config)
self._t2r_model = t2r_model
# The location of the last checkpoint loaded.
self._current_checkpoint_path = None
self._tf_global_step = tf.train.get_or_create_global_step()
# The PREDICT graph is generated which contains only the model specific
# variables and not training specific variables, e.g. Adam, Momentum.
var_list = contrib_framework.get_variables()
self._saver = tf.train.Saver(var_list=var_list)
# Default init op in case init_randomly is called.
self._global_init_op = tf.global_variables_initializer()
self._model_was_restored = False
