def test_cast_bfloat16_raises(self):
# Note, this preprocessor will alter the input data type from float32 to
# int32 which will trigger ValueError when attempting to cast.
preprocessor = MockPreprocessor(
model_feature_specification_fn=lambda mode: _FEATURE_SPEC_CAST,
model_label_specification_fn=lambda mode: _LABEL_SPEC_CAST)
tpu_preprocessor = tpu_preprocessor_wrapper.TPUPreprocessorWrapper(
preprocessor=preprocessor)
out_feature_spec = preprocessor.get_out_feature_specification(_MODE_TRAIN)
del out_feature_spec['optional_value']
out_label_spec = preprocessor.get_out_label_specification(_MODE_TRAIN)
del out_label_spec['optional_value']
self.assertDictEqual(
tpu_preprocessor.get_out_feature_specification(_MODE_TRAIN),
out_feature_spec)
self.assertDictEqual(
tpu_preprocessor.get_out_label_specification(_MODE_TRAIN),
out_label_spec)
features = tensorspec_utils.make_placeholders(
tpu_preprocessor.get_in_feature_specification(_MODE_TRAIN),
batch_size=2)
labels = tensorspec_utils.make_placeholders(
tpu_preprocessor.get_in_label_specification(_MODE_TRAIN), batch_size=2)
# Make sure the exception gets triggered if features and labels are passed.
with self.assertRaises(ValueError):
tpu_preprocessor.preprocess(
features=features, labels=labels, mode=_MODE_TRAIN)
# Make sure the exception gets triggered if only features are passed.
with self.assertRaises(ValueError):
tpu_preprocessor.preprocess(
features=features, labels=None, mode=_MODE_TRAIN)
def test_pack_flat_sequence_to_spec_structure(self):
subset_placeholders = utils.make_placeholders(mock_nested_subset_spec)
flattened_subset_placeholders = utils.flatten_spec_structure(
subset_placeholders)
packed_subset_placeholders = utils.pack_flat_sequence_to_spec_structure(
mock_nested_subset_spec, flattened_subset_placeholders)
utils.assert_equal(subset_placeholders, packed_subset_placeholders)
utils.assert_equal(mock_nested_subset_spec,
packed_subset_placeholders,
ignore_batch=True)
placeholders = utils.make_placeholders(mock_nested_spec)
flattened_placeholders = utils.flatten_spec_structure(placeholders)
packed_placeholders = utils.pack_flat_sequence_to_spec_structure(
mock_nested_subset_spec, flattened_placeholders)
# We only subselect what we need in pack_flat_sequence_to_spec_structure,
# hence, we should recover what we wanted.
utils.assert_equal(mock_nested_subset_spec,
packed_placeholders,
ignore_batch=True)
utils.assert_equal(subset_placeholders, packed_placeholders)
packed_optional_placeholders = utils.pack_flat_sequence_to_spec_structure(
mock_nested_optional_spec, flattened_placeholders)
# Although mock_nested_optional_spec would like more tensors
# flattened_placeholders cannot provide them, fortunately they are optional.
utils.assert_required(packed_optional_placeholders, placeholders)
utils.assert_required(mock_nested_spec,
packed_optional_placeholders,
ignore_batch=True)
def test_cast_bfloat16_success(self):
preprocessor = noop_preprocessor.NoOpPreprocessor(
model_feature_specification_fn=lambda mode: _FEATURE_SPEC_CAST,
model_label_specification_fn=lambda mode: _LABEL_SPEC_CAST)
tpu_preprocessor = tpu_preprocessor_wrapper.TPUPreprocessorWrapper(
preprocessor=preprocessor)
# The spec structure elements with bfloat16 are converted to float32 within
# the TPUPreprocessorWrapper such that we can create proper parser and
# do CPU preprocessing.
feature_spec = preprocessor.get_in_feature_specification(_MODE_TRAIN)
feature_spec.data_bfloat16 = tensorspec_utils.ExtendedTensorSpec.from_spec(
spec=feature_spec.data_bfloat16, dtype=tf.float32)
label_spec = preprocessor.get_in_label_specification(_MODE_TRAIN)
label_spec.optional_value = tensorspec_utils.ExtendedTensorSpec.from_spec(
spec=label_spec.optional_value, dtype=tf.float32)
self.assertDictEqual(
tpu_preprocessor.get_in_feature_specification(_MODE_TRAIN),
feature_spec)
self.assertDictEqual(
tpu_preprocessor.get_in_label_specification(_MODE_TRAIN),
label_spec)
out_feature_spec = preprocessor.get_out_feature_specification(
_MODE_TRAIN)
del out_feature_spec['optional_value']
out_label_spec = preprocessor.get_out_label_specification(_MODE_TRAIN)
del out_label_spec['optional_value']
self.assertDictEqual(
tpu_preprocessor.get_out_feature_specification(_MODE_TRAIN),
out_feature_spec)
self.assertDictEqual(
tpu_preprocessor.get_out_label_specification(_MODE_TRAIN),
out_label_spec)
features = tensorspec_utils.make_placeholders(
tpu_preprocessor.get_in_feature_specification(_MODE_TRAIN),
batch_size=2)
labels = tensorspec_utils.make_placeholders(
tpu_preprocessor.get_in_label_specification(_MODE_TRAIN),
batch_size=2)
# Make sure features and labels are transformed correctly. Basically
# float32 is replaced with bfloat16 for the specs which ask for bfloat16.
out_features, out_labels = tpu_preprocessor.preprocess(
features=features, labels=labels, mode=_MODE_TRAIN)
for ref_key, ref_value in out_features.items():
self.assertEqual(out_features[ref_key].dtype, ref_value.dtype)
for ref_key, ref_value in out_labels.items():
self.assertEqual(out_labels[ref_key].dtype, ref_value.dtype)
# Make sure features without labels are transformed correctly. Basically
# float32 is replaced with bfloat16 for the specs which ask for bfloat16.
out_features, out_labels = tpu_preprocessor.preprocess(
features=features, labels=None, mode=_MODE_TRAIN)
self.assertIsNone(out_labels)
for ref_key, ref_value in out_features.items():
self.assertEqual(out_features[ref_key].dtype, ref_value.dtype)
def test_make_placeholders(self, collection_type):
spec = self._make_tensorspec_collection(collection_type)
placeholders = utils.make_placeholders(spec)
placeholder_spec = utils.tensorspec_from_tensors(placeholders)
utils.assert_equal(spec, placeholder_spec, ignore_batch=True)
with self.assertRaises(ValueError):
utils.assert_equal(spec, placeholder_spec, ignore_batch=False)
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._get_input_features_for_receiver_fn())
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 (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 infer_base_model_output_dtypes(self, mode, params):
"""Infer the dtypes of the model in a separate graph.
Args:
mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
params: An optional dict of hyper parameters that will be passed into
input_fn and model_fn. Keys are names of parameters, values are basic
python types. There are reserved keys for TPUEstimator, including
'batch_size'.
Returns:
dtypes: A dict containing all output dtypes {str: dtype}.
"""
dtype_inference_graph = tf.Graph()
with dtype_inference_graph.as_default():
with tf.variable_scope('IGNORE_ONLY_TO_INFER_OUTPUT_DTYPES'):
# In this graph we can now create placeholders in order to infer the
# right dtype of the outputs.
feature_spec = self.get_feature_specification(mode)
features_dtype = utils.make_placeholders(
feature_spec.condition.features, batch_size=-1)
labels_dtype = utils.make_placeholders(
feature_spec.condition.labels, batch_size=-1)
# We need to infer the output dtypes.
features_condition = utils.flatten_spec_structure(
features_dtype)
labels_condition = utils.flatten_spec_structure(labels_dtype)
infered_outputs = self._base_model.inference_network_fn(
features=features_condition,
labels=labels_condition,
mode=mode,
params=params)
dtypes = {}
for key, value in infered_outputs.items():
dtypes[key] = value.dtype
return dtypes
def test_pack_flat_sequence_to_spec_structure_ensure_order(self):
test_spec = utils.TensorSpecStruct()
test_spec.b = utils.ExtendedTensorSpec(
shape=(1,), dtype=tf.float32, name='b')
test_spec.a = utils.ExtendedTensorSpec(
shape=(1,), dtype=tf.float32, name='a')
test_spec.c = utils.ExtendedTensorSpec(
shape=(1,), dtype=tf.float32, name='c')
placeholders = utils.make_placeholders(test_spec)
packed_placeholders = utils.pack_flat_sequence_to_spec_structure(
test_spec, placeholders)
for pos, order_name in enumerate(['a', 'b', 'c']):
self.assertEqual(list(packed_placeholders.keys())[pos], order_name)
self.assertEqual(
list(packed_placeholders.values())[pos].op.name, order_name)
def restore(self):
# Forcing both session and graph to be defaults here to force a graph
# context if this ever gets used during Eager execution. Makes testing
# easier too.
with self._session.as_default(), self._graph.as_default():
if not super(SavedModelTF1Predictor, self).restore():
return False
self._features = tensorspec_utils.make_placeholders(
self._feature_spec, batch_size=None)
self._predictions = super(SavedModelTF1Predictor,
self).predict(self._features)
# After loading the model we need to make sure we initialize the
# variables.
variables = (tf.compat.v1.global_variables() +
tf.compat.v1.local_variables())
self._session.run(tf.compat.v1.variables_initializer(variables))
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 test_validate_flatten_and_pack(self):
# An example data pipeline.
# Some input generator creates input features according to some spec.
input_features = utils.make_placeholders(mock_nested_spec)
# Assume a preprocessor has altered these input_features and we want
# to pass the data on to the next stage, then we simply assure that
# our output is according to our spec and flatten.
flat_input_features = utils.validate_and_flatten(
mock_nested_optional_spec, input_features, ignore_batch=True)
utils.assert_required(
mock_nested_optional_spec, input_features, ignore_batch=True)
# Then e.g. the model_fn receives the flat_input_spec and validates
# that it is according to it's requirements and packs it back into the
# spec structure.
output_features = utils.validate_and_pack(
mock_nested_subset_spec, flat_input_features, ignore_batch=True)
utils.assert_required(
mock_nested_subset_spec, output_features, ignore_batch=True)
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
def preprocess(preprocessor, feature_spec, label_spec, flatten=False):
with tf.Session() as sess:
feature_placeholders = tensorspec_utils.make_placeholders(
feature_spec, batch_size=1)
label_placeholders = None
if label_spec is not None:
label_placeholders = tensorspec_utils.make_placeholders(
label_spec, batch_size=1)
# Normally we want our features and labels to be flattened.
# However we support not flattened hierarchies as well.
if flatten:
feature_placeholders = tensorspec_utils.flatten_spec_structure(
feature_placeholders)
if label_spec is not None:
label_placeholders = tensorspec_utils.flatten_spec_structure(
label_placeholders)
(features_preprocessed, labels_preprocessed) = preprocessor.preprocess(
features=feature_placeholders,
labels=label_placeholders,
mode=tf.estimator.ModeKeys.TRAIN)
# We create a mapping of {key: np.array} or a namedtuple spec structure.
np_feature_spec = tensorspec_utils.make_random_numpy(
feature_spec, batch_size=1)
if label_placeholders is not None:
np_label_spec = tensorspec_utils.make_random_numpy(
label_spec, batch_size=1)
# We create our feed dict which basically consists of
# {placeholders: np.array}.
feed_dict = tensorspec_utils.map_feed_dict(feature_placeholders,
np_feature_spec,
ignore_batch=True)
if label_placeholders is not None:
feed_dict = tensorspec_utils.map_feed_dict(label_placeholders,
np_label_spec,
feed_dict,
ignore_batch=True)
fetch_results = [features_preprocessed]
if label_placeholders is not None:
fetch_results.append(labels_preprocessed)
np_preprocessed = sess.run(
fetch_results, feed_dict=feed_dict)
np_features_preprocessed = np_preprocessed[0]
if label_placeholders is not None:
np_labels_preprocessed = np_preprocessed[1]
np_feature_spec = tensorspec_utils.flatten_spec_structure(
np_feature_spec)
if label_placeholders is not None:
np_label_spec = tensorspec_utils.flatten_spec_structure(np_label_spec)
for key, value in np_feature_spec.items():
np.testing.assert_allclose(value, np_features_preprocessed[key])
if label_placeholders is not None:
for key, value in np_label_spec.items():
np.testing.assert_allclose(value, np_labels_preprocessed[key])
