def ReadExamples(path):
print("Reading", path)
records = list(tf.python_io.tf_record_iterator(path, TF_RECORD_CONFIG))
num_records = len(records)
features = {
'x': tf.FixedLenFeature([], tf.string),
'pi': tf.FixedLenFeature([], tf.string),
'outcome': tf.FixedLenFeature([], tf.float32),
'n': tf.FixedLenFeature([], tf.int64, default_value=[-1]),
'q': tf.FixedLenFeature([], tf.float32, default_value=[-1]),
'c': tf.FixedLenFeature([], tf.int64, default_value=[-1]),
}
parsed = tf.parse_example(records, features)
x = tf.decode_raw(parsed['x'], tf.uint8)
x = tf.cast(x, tf.float32)
x = tf.reshape(x, [num_records, go.N, go.N, -1])
x = x.eval()
pi = tf.decode_raw(parsed['pi'], tf.float32)
pi = tf.reshape(pi, [num_records, go.N * go.N + 1])
pi = pi.eval()
outcome = parsed['outcome'].eval()
n = parsed['n'].eval()
q = parsed['q'].eval()
c = parsed['c'].eval()
return [ParsedExample(*args) for args in zip(x, pi, outcome, q, n, c)]
def parse_fn(example):
feature_map = {
"example_id": tf.FixedLenFeature([], tf.string),
"x": tf.FixedLenFeature([28 * 28 // 2], tf.float32),
}
features = tf.parse_example(example, features=feature_map)
return features, {}
def parse_fn(example):
feature_map = dict()
feature_map['fids'] = tf.VarLenFeature(tf.int64)
feature_map['example_id'] = tf.FixedLenFeature([], tf.string)
feature_map["y"] = tf.FixedLenFeature([], tf.int64)
features = tf.parse_example(example, features=feature_map)
return features, dict(y=features.pop('y'))
def _parser(serialized_example):
"""Parses a single tf.Example into image and label tensors."""
features = tf.parse_example(
[serialized_example],
features={
'image/encoded': tf.VarLenFeature(dtype=tf.float32),
'image/segmentation/mask':
tf.VarLenFeature(dtype=tf.float32),
})
image = features['image/encoded']
if isinstance(image, tf.SparseTensor):
image = tf.sparse_tensor_to_dense(image)
gt_mask = features['image/segmentation/mask']
if isinstance(gt_mask, tf.SparseTensor):
gt_mask = tf.sparse_tensor_to_dense(gt_mask)
image_size, label_size = self.get_input_shapes(params)
image = tf.reshape(image, image_size)
gt_mask = tf.reshape(gt_mask, label_size)
if params.use_bfloat16:
image = tf.cast(image, dtype=tf.bfloat16)
gt_mask = tf.cast(gt_mask, dtype=tf.bfloat16)
logging.info('debug input %s %s', image, gt_mask)
return image, gt_mask
def parse_fn(example):
feature_map = dict()
feature_map['example_id'] = tf.FixedLenFeature([], tf.string)
feature_map['x'] = tf.FixedLenFeature([28 * 28 // 2], tf.float32)
feature_map['y'] = tf.FixedLenFeature([], tf.int64)
features = tf.parse_example(example, features=feature_map)
return features, dict(y=features.pop('y'))
def _serving_input_receiver_fn(
long_seq_length: int,
global_seq_length: int,
) -> tf.estimator.export.ServingInputReceiver:
"""Creates an input function to parse input features for inference.
This function defines format of the inputs to the exported HotpotQA model.
at inference time.
Args:
long_seq_length: The long input len.
global_seq_length: The global input len.
Returns:
The ServingInputReceiver fn.
"""
# An input receiver that expects a vector of serialized `tf.Example`s.
serialized_tf_example = tf.placeholder(dtype=tf.string,
shape=[None],
name="serialized_tf_example")
receiver_tensors = {"serialized_tf_example": serialized_tf_example}
schema = run_finetuning_lib.get_inference_name_to_features(
long_seq_length, global_seq_length)
features = tf.parse_example(serialized_tf_example, schema)
return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
def parse_example_batch(examples, symbolic_properties):
"""Parses a batch of tf.Examples.
Args:
examples: String tensor with shape [batch_size]. A batch of serialized
tf.Example protos.
symbolic_properties: List of strings, symbolic properties to concatenate on
embedding as conditions. Those symbolic_properties will be read from
input data.
Returns:
A feature dict. It contains key 'expression_string' with a string tensor of
expressions with shape [batch_size]. It also contain keys in
symbolic_properties with float tensors.
"""
features_to_extract = {
'expression_string': tf.FixedLenFeature([], tf.string),
}
for symbolic_property in symbolic_properties:
features_to_extract[symbolic_property] = tf.FixedLenFeature([],
tf.float32)
features = tf.parse_example(examples, features=features_to_extract)
return features
def parse_fn(example):
feature_map = {}
feature_map['example_id'] = tf.FixedLenFeature([], tf.string)
feature_map['y'] = tf.FixedLenFeature([], tf.int64)
features = tf.parse_example(example, features=feature_map)
labels = {'y': features.pop('y')}
return features, labels
def batch_parse_tf_example(batch_size, layout, example_batch):
"""
Args:
batch_size: batch size
layout: 'nchw' or 'nhwc'
example_batch: a batch of tf.Example
Returns:
A tuple (feature_tensor, dict of output tensors)
"""
planes = dual_net.get_features_planes()
features = {
'x': tf.FixedLenFeature([], tf.string),
'pi': tf.FixedLenFeature([], tf.string),
'outcome': tf.FixedLenFeature([], tf.float32),
}
parsed = tf.parse_example(example_batch, features)
x = tf.decode_raw(parsed['x'], tf.uint8)
x = tf.cast(x, tf.float32)
if layout == 'nhwc':
shape = [batch_size, go.N, go.N, planes]
else:
shape = [batch_size, planes, go.N, go.N]
x = tf.reshape(x, shape)
pi = tf.decode_raw(parsed['pi'], tf.float32)
pi = tf.reshape(pi, [batch_size, go.N * go.N + 1])
outcome = parsed['outcome']
outcome.set_shape([batch_size])
return x, {'pi_tensor': pi, 'value_tensor': outcome}
def build_graph(parameters):
"""Build the graph for parse_example tests."""
feature_dtype = parameters["feature_dtype"]
feature_shape = parameters["feature_shape"]
is_dense = parameters["is_dense"]
input_value = tf.compat.v1.placeholder(dtype=tf.string,
name="input",
shape=[1])
if is_dense:
feature_default_value = np.zeros(shape=feature_shape)
if feature_dtype == tf.string:
feature_default_value = np.array(["missing"] *
feature_shape[0])
features = {
"x":
tf.FixedLenFeature(shape=feature_shape,
dtype=feature_dtype,
default_value=feature_default_value)
}
else: # Sparse
features = {"x": tf.VarLenFeature(dtype=feature_dtype)}
out = tf.parse_example(input_value, features)
output_tensor = out["x"]
if not is_dense:
output_tensor = out["x"].values
return [input_value], [output_tensor]
def _serving_input_receiver_fn():
"""Creates an input function for serving."""
seq_len = FLAGS.max_seq_length
serialized_example = tf.placeholder(dtype=tf.string,
shape=[None],
name="serialized_example")
features = {
"input_ids": tf.FixedLenFeature([seq_len], dtype=tf.int64),
"input_mask": tf.FixedLenFeature([seq_len], dtype=tf.int64),
"segment_ids": tf.FixedLenFeature([seq_len], dtype=tf.int64),
}
feature_map = tf.parse_example(serialized_example, features=features)
feature_map["is_real_example"] = tf.constant(1, dtype=tf.int32)
feature_map["label_ids"] = tf.constant(0, dtype=tf.int32)
# tf.Example only supports tf.int64, but the TPU only supports tf.int32.
# So cast all int64 to int32.
for name in feature_map.keys():
t = feature_map[name]
if t.dtype == tf.int64:
t = tf.to_int32(t)
feature_map[name] = t
return tf.estimator.export.ServingInputReceiver(
features=feature_map, receiver_tensors=serialized_example)
def parse_mnist_tfrecord(serialized_example: tf.Tensor) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
"""
Parse a TFRecord representing a single MNIST data point into an input
feature tensor and a label tensor.
Returns: (features: Dict[str, Tensor], label: Tensor)
"""
raw = tf.parse_example(
serialized=serialized_example, features={"image_raw": tf.FixedLenFeature([], tf.string)}
)
image = tf.decode_raw(raw["image_raw"], tf.float32)
label_dict = tf.parse_example(
serialized=serialized_example, features={"label": tf.FixedLenFeature(1, tf.int64)}
)
return {"image": image}, label_dict["label"]
def _generate_tflite_for_parse_example_with_string(export_dir):
"""Generates TFLite flatbuffer for parse example with string.
Args:
export_dir: The directory to which the flatbuffer should be written.
"""
with tf.Session(
graph=tf.Graph(),
config=tf.ConfigProto(log_device_placement=True)) as sess:
serialized_tf_example = tf.placeholder(
tf.string, name="input", shape=[None])
tf_example = tf.parse_example(serialized_tf_example,
_get_feature_spec())
converter = tf.lite.TFLiteConverter.from_session(
sess, [serialized_tf_example], [tf_example["x"], tf_example["y"]])
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
tflite_model = converter.convert()
predict_signature_def = _build_predict_signature(serialized_tf_example,
tf_example["x"],
tf_example["y"])
k = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
tflite_model = signature_def_utils.set_signature_defs(
tflite_model, {k: predict_signature_def})
open(export_dir + "/model.tflite", "wb").write(tflite_model)
def _serving_input_fn():
"""An input_fn that expects a serialized tf.Example."""
serialized_example = tf.placeholder(
dtype=tf.string, shape=[None], name="examples")
receiver_tensors = {"examples": serialized_example}
name_to_features = {
"input_ids_1": tf.FixedLenFeature([max_seq_length], tf.int64),
"input_mask_1": tf.FixedLenFeature([max_seq_length], tf.int64),
"input_ids_2": tf.FixedLenFeature([max_seq_length], tf.int64),
"input_mask_2": tf.FixedLenFeature([max_seq_length], tf.int64),
"documents_match_labels": tf.FixedLenFeature([1], tf.float32, 0)
}
if (FLAGS.train_mode == constants.TRAIN_MODE_PRETRAIN or
FLAGS.train_mode == constants.TRAIN_MODE_JOINT_TRAIN):
# This is to support model export during model pretraining or
# joint-training process.
name_to_features["masked_lm_positions_1"] = tf.FixedLenFeature(
[max_predictions_per_seq], tf.int64)
name_to_features["masked_lm_ids_1"] = tf.FixedLenFeature(
[max_predictions_per_seq], tf.int64)
name_to_features["masked_lm_weights_1"] = tf.FixedLenFeature(
[max_predictions_per_seq], tf.float32)
name_to_features["masked_lm_positions_2"] = tf.FixedLenFeature(
[max_predictions_per_seq], tf.int64)
name_to_features["masked_lm_ids_2"] = tf.FixedLenFeature(
[max_predictions_per_seq], tf.int64)
name_to_features["masked_lm_weights_2"] = tf.FixedLenFeature(
[max_predictions_per_seq], tf.float32)
parsed_features = tf.parse_example(serialized_example, name_to_features)
# As tf.Example only supports tf.int64, but the TPU only supports
# tf.int32, we need to cast all int64 to int32.
parsed_features["input_ids_1"] = tf.cast(parsed_features["input_ids_1"],
tf.int32)
parsed_features["input_ids_2"] = tf.cast(parsed_features["input_ids_2"],
tf.int32)
parsed_features["documents_match_labels"] = tf.cast(
parsed_features["documents_match_labels"], tf.float32)
parsed_features["input_mask_1"] = tf.cast(parsed_features["input_mask_1"],
tf.int32)
parsed_features["input_mask_2"] = tf.cast(parsed_features["input_mask_2"],
tf.int32)
if (FLAGS.train_mode == constants.TRAIN_MODE_PRETRAIN or
FLAGS.train_mode == constants.TRAIN_MODE_JOINT_TRAIN):
parsed_features["masked_lm_ids_1"] = tf.cast(
parsed_features["masked_lm_ids_1"], tf.int32)
parsed_features["masked_lm_ids_2"] = tf.cast(
parsed_features["masked_lm_ids_2"], tf.int32)
parsed_features["masked_lm_weights_1"] = tf.cast(
parsed_features["masked_lm_weights_1"], tf.float32)
parsed_features["masked_lm_weights_2"] = tf.cast(
parsed_features["masked_lm_weights_2"], tf.float32)
parsed_features["masked_lm_positions_1"] = tf.cast(
parsed_features["masked_lm_positions_1"], tf.int32)
parsed_features["masked_lm_positions_2"] = tf.cast(
parsed_features["masked_lm_positions_2"], tf.int32)
return tf.estimator.export.ServingInputReceiver(
features=parsed_features, receiver_tensors=receiver_tensors)
def serving_input_receiver_fn():
feature_map = {"fids": tf.VarLenFeature(tf.int64)}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
record_batch = tf.placeholder(dtype=tf.string, name='examples')
features = tf.parse_example(record_batch, features=feature_map)
return tf.estimator.export.ServingInputReceiver(
features, {'examples': record_batch})
def parse_fn(example):
feature_map = {
"x_{0}".format(i): tf.VarLenFeature(tf.int64)
for i in range(512)
}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
features = tf.parse_example(example, features=feature_map)
labels = {}
return features, labels
def input_fn(bridge, trainer_master):
dataset = flt.data.DataBlockLoader(256, 'follower', bridge, trainer_master)
feature_map = {
"example_id": tf.FixedLenFeature([], tf.string),
"x": tf.FixedLenFeature([28 * 28 // 2], tf.float32),
}
record_batch = dataset.make_batch_iterator().get_next()
features = tf.parse_example(record_batch, features=feature_map)
return features, {}
def parse_fn(example):
feature_map = {}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
feature_map['fids'] = tf.VarLenFeature(tf.int64)
# feature_map['y'] = tf.FixedLenFeature([], tf.int64)
features = tf.parse_example(example, features=feature_map)
# labels = {'y': features.pop('y')}
labels = {'y': tf.constant(0)}
return features, labels
def serving_input_receiver_fn():
feature_map = {
"example_id": tf.FixedLenFeature([], tf.string),
"x": tf.FixedLenFeature([28 * 28 // 2], tf.float32),
}
record_batch = tf.placeholder(dtype=tf.string, name='examples')
features = tf.parse_example(record_batch, features=feature_map)
return tf.estimator.export.ServingInputReceiver(features,
{'examples': record_batch})
def parse_fn(example):
feature_map = {
"x_{0}".format(i): tf.VarLenFeature(tf.int64)
for i in range(512)
}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
feature_map["y"] = tf.FixedLenFeature([], tf.int64)
features = tf.parse_example(example, features=feature_map)
return features, dict(y=features.pop('y'))
def input_fn(bridge, trainer_master=None):
dataset = flt.data.DataBlockLoader(
args.batch_size, ROLE, bridge, trainer_master)
feature_map = {"fids": tf.VarLenFeature(tf.int64)}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
record_batch = dataset.make_batch_iterator().get_next()
features = tf.parse_example(record_batch, features=feature_map)
return features, None
def parse_fn(example):
feature_map = {
"example_id": tf.FixedLenFeature([], tf.string),
"x": tf.FixedLenFeature([28 * 28 // 2], tf.float32),
"y": tf.FixedLenFeature([], tf.int64)
}
features = tf.parse_example(example, features=feature_map)
labels = {'y': features.pop('y')}
print("Leader worker input_fn ", labels)
return features, labels
def input_fn(bridge, trainer_master=None):
dataset = flt.data.DataBlockLoader(
args.batch_size, ROLE, bridge, trainer_master)
feature_map = {"x_{0}".format(i): tf.VarLenFeature(
tf.int64) for i in range(512)}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
feature_map["y"] = tf.FixedLenFeature([], tf.int64)
record_batch = dataset.make_batch_iterator().get_next()
features = tf.parse_example(record_batch, features=feature_map)
labels = {'y': features.pop('y')}
return features, labels
def serving_input_receiver_fn():
feature_map = {
"x_{0}".format(i): tf.VarLenFeature(tf.int64)
for i in range(512)
}
feature_map["example_id"] = tf.FixedLenFeature([], tf.string)
record_batch = tf.placeholder(dtype=tf.string, name='examples')
features = tf.parse_example(record_batch, features=feature_map)
features['act1_f'] = tf.placeholder(dtype=tf.float32, name='act1_f')
receiver_tensors = {'examples': record_batch, 'act1_f': features['act1_f']}
return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
def serving_input_receiver_fn():
feature_spec = {
"unique_ids": tf.FixedLenFeature([], tf.int64),
"input_ids": tf.FixedLenFeature([FLAGS.max_seq_length],
tf.int64),
"input_mask": tf.FixedLenFeature([FLAGS.max_seq_length],
tf.int64),
"segment_ids": tf.FixedLenFeature([FLAGS.max_seq_length],
tf.int64),
"p_mask": tf.FixedLenFeature([FLAGS.max_seq_length], tf.int64)
}
serialized_tf_example = tf.placeholder(
dtype=tf.string,
shape=[FLAGS.predict_batch_size],
name='input_example_tensor')
receiver_tensors = {'examples': serialized_tf_example}
features = tf.parse_example(serialized_tf_example, feature_spec)
return tf.estimator.export.ServingInputReceiver(
features, receiver_tensors)
def _eval_input_receiver_fn(tf_transform_output, schema):
"""Build everything needed for the tf-model-analysis to run the model.
Args:
tf_transform_output: A TFTransformOutput.
schema: the schema of the input data.
Returns:
EvalInputReceiver function, which contains:
- Tensorflow graph which parses raw untransformed features, applies the
tf-transform preprocessing operators.
- Set of raw, untransformed features.
- Label against which predictions will be compared.
"""
# Notice that the inputs are raw features, not transformed features here.
raw_feature_spec = _get_raw_feature_spec(schema)
serialized_tf_example = tf.placeholder(dtype=tf.string,
shape=[None],
name='input_example_tensor')
# Add a parse_example operator to the tensorflow graph, which will parse
# raw, untransformed, tf examples.
features = tf.parse_example(serialized_tf_example, raw_feature_spec)
# Now that we have our raw examples, process them through the tf-transform
# function computed during the preprocessing step.
transformed_features = tf_transform_output.transform_raw_features(features)
# The key name MUST be 'examples'.
receiver_tensors = {'examples': serialized_tf_example}
# NOTE: Model is driven by transformed features (since training works on the
# materialized output of TFT, but slicing will happen on raw features.
features.update(transformed_features)
return tfma.export.EvalInputReceiver(
features=features,
receiver_tensors=receiver_tensors,
labels=transformed_features[_transformed_name(_LABEL_KEY)])
def serving_input_receiver_fn(
model_config: modeling.EtcConfig,
long_seq_len: int,
global_seq_len: int,
candidate_ignore_hard_g2l: bool = True,
query_ignore_hard_g2l: bool = True,
enable_l2g_linking: bool = True
) -> tf.estimator.export.ServingInputReceiver:
"""Creates an input function to parse input features for inference."""
# An input receiver that expects a vector of serialized `tf.Example`s.
serialized_tf_example = tf.placeholder(dtype=tf.string,
shape=[None],
name="serialized_tf_example")
receiver_tensors = {"serialized_tf_example": serialized_tf_example}
schema = get_model_schema(long_seq_len, global_seq_len)
features = tf.parse_example(serialized_tf_example, schema)
features = _add_side_input_features(
features=features,
model_config=model_config,
candidate_ignore_hard_g2l=candidate_ignore_hard_g2l,
query_ignore_hard_g2l=query_ignore_hard_g2l,
enable_l2g_linking=enable_l2g_linking)
return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
def _generate_saved_model_for_half_plus_two(export_dir,
as_text=False,
as_tflite=False,
use_main_op=False,
include_mlmd=False,
device_type="cpu"):
"""Generates SavedModel for half plus two.
Args:
export_dir: The directory to which the SavedModel should be written.
as_text: Writes the SavedModel protocol buffer in text format to disk.
as_tflite: Writes the Model in Tensorflow Lite format to disk.
use_main_op: Whether to supply a main op during SavedModel build time.
include_mlmd: Whether to include an MLMD key in the SavedModel.
device_type: Device to force ops to run on.
"""
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
device_name = "/cpu:0"
if device_type == "gpu":
device_name = "/gpu:0"
with tf.Session(graph=tf.Graph(),
config=tf.ConfigProto(log_device_placement=True)) as sess:
with tf.device(device_name):
# Set up the model parameters as variables to exercise variable loading
# functionality upon restore.
a = tf.Variable(0.5, name="a")
b = tf.Variable(2.0, name="b")
c = tf.Variable(3.0, name="c")
# Create a placeholder for serialized tensorflow.Example messages to be
# fed.
serialized_tf_example = tf.placeholder(tf.string,
name="tf_example",
shape=[None])
# Parse the tensorflow.Example looking for a feature named "x" with a
# single floating point value.
feature_configs = {
"x":
tf.FixedLenFeature([1], dtype=tf.float32),
"x2":
tf.FixedLenFeature([1], dtype=tf.float32, default_value=[0.0])
}
# parse_example only works on CPU
with tf.device("/cpu:0"):
tf_example = tf.parse_example(serialized_tf_example,
feature_configs)
# Use tf.identity() to assign name
x = tf.identity(tf_example["x"], name="x")
if device_type == "mkl":
# Create a small convolution op to trigger MKL
# The op will return 0s so this won't affect the
# resulting calculation.
o1 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros(
(1, 16, 16, 1)))
y = o1[0, 0, 0, 0] + tf.add(tf.multiply(a, x), b)
else:
y = tf.add(tf.multiply(a, x), b)
y = tf.identity(y, name="y")
if device_type == "mkl":
# Create a small convolution op to trigger MKL
# The op will return 0s so this won't affect the
# resulting calculation.
o2 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros(
(1, 16, 16, 1)))
y2 = o2[0, 0, 0, 0] + tf.add(tf.multiply(a, x), c)
else:
y2 = tf.add(tf.multiply(a, x), c)
y2 = tf.identity(y2, name="y2")
x2 = tf.identity(tf_example["x2"], name="x2")
if device_type == "mkl":
# Create a small convolution op to trigger MKL
# The op will return 0s so this won't affect the
# resulting calculation.
o3 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros(
(1, 16, 16, 1)))
y3 = o3[0, 0, 0, 0] + tf.add(tf.multiply(a, x2), c)
else:
# Add separate constants for x2, to prevent optimizers like TF-TRT from
# fusing the paths to compute y/y2 and y3 together.
a2 = tf.Variable(0.5, name="a2")
c2 = tf.Variable(3.0, name="c2")
y3 = tf.add(tf.multiply(a2, x2), c2)
y3 = tf.identity(y3, name="y3")
assign_filename_op = _create_asset_file()
# Set up the signature for Predict with input and output tensor
# specification.
predict_input_tensor = tf.saved_model.utils.build_tensor_info(x)
predict_signature_inputs = {"x": predict_input_tensor}
predict_output_tensor = tf.saved_model.utils.build_tensor_info(y)
predict_signature_outputs = {"y": predict_output_tensor}
predict_signature_def = (
tf.saved_model.signature_def_utils.build_signature_def(
predict_signature_inputs, predict_signature_outputs,
tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
signature_def_map = {
"regress_x_to_y":
_build_regression_signature(serialized_tf_example, y),
"regress_x_to_y2":
_build_regression_signature(serialized_tf_example, y2),
"regress_x2_to_y3":
_build_regression_signature(x2, y3),
"classify_x_to_y":
_build_classification_signature(serialized_tf_example, y),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
# Initialize all variables and then save the SavedModel.
sess.run(tf.global_variables_initializer())
if as_tflite:
converter = tf.lite.TFLiteConverter.from_session(sess, [x], [y])
tflite_model = converter.convert()
open(export_dir + "/model.tflite", "wb").write(tflite_model)
else:
if use_main_op:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(
tf.GraphKeys.ASSET_FILEPATHS),
main_op=tf.group(tf.saved_model.main_op.main_op(),
assign_filename_op))
else:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(
tf.GraphKeys.ASSET_FILEPATHS),
main_op=tf.group(assign_filename_op))
if not as_tflite:
builder.save(as_text)
if include_mlmd:
_write_mlmd(export_dir, "test_mlmd_uuid")
def parse_wrapper(example, spec_dict):
"""Wrap tf.parse_example to support bfloat16 dtypes.
This allows models which declare bfloat16 as inputs to not require an
additional preprocessing step to cast all inputs from float32 to bfloat16.
Consider this to be analogous to JPEG decoding in the data step.
Args:
example: TFExample
spec_dict: Dictionary of feature name -> tf.FixedLenFeature
Returns:
Parsed feature map
"""
def is_bfloat_feature(value):
return value.dtype == tf.bfloat16
def maybe_map_bfloat(value):
"""Maps bfloat16 to float32."""
if is_bfloat_feature(value):
if isinstance(value, tf.FixedLenFeature):
return tf.FixedLenFeature(
value.shape, tf.float32, default_value=value.default_value)
elif isinstance(value, tf.VarLenFeature):
return tf.VarLenFeature(
value.shape, tf.float32, default_value=value.default_value)
else:
return tf.FixedLenSequenceFeature(
value.shape, tf.float32, default_value=value.default_value)
return value
# Change bfloat features to float32 for parsing.
new_spec_dict = {
k: maybe_map_bfloat(v) for k, v in six.iteritems(spec_dict)
}
for k, v in six.iteritems(new_spec_dict):
if v.dtype not in [tf.float32, tf.string, tf.int64]:
raise ValueError('Feature specification with invalid data type for '
'tf.Example parsing: "%s": %s' % (k, v.dtype))
# Separate new_spec_dict into Context and Sequence features. In the event
# that there are no SequenceFeatures, the context_features dictionary
# (containing FixedLenFeatures) is passed to tf.parse_examples.
context_features, sequence_features = {}, {}
for k, v in six.iteritems(new_spec_dict):
v = maybe_map_bfloat(v)
if isinstance(v, tf.FixedLenSequenceFeature):
sequence_features[k] = v
elif isinstance(v, tf.FixedLenFeature):
context_features[k] = v
elif isinstance(v, tf.VarLenFeature):
context_features[k] = v
else:
raise ValueError(
'Only FixedLenFeature and FixedLenSequenceFeature are currently '
'supported.')
# If there are any sequence features, we use parse_sequence_example.
if sequence_features:
# Filter out '_length' context features; don't parse them from records.
for parse_name in sequence_features:
# Sometimes, the '_length' context feature doesn't exist.
if parse_name + '_length' in context_features:
del context_features[parse_name + '_length']
result, sequence_result, feature_lengths = tf.io.parse_sequence_example(
example,
context_features=context_features,
sequence_features=sequence_features)
result.update(sequence_result)
# Augment the parsed tensors with feature length tensors.
for parse_name, length_tensor in feature_lengths.items():
result[parse_name + '_length'] = length_tensor
else:
result = tf.parse_example(example, context_features)
to_convert = [
k for k, v in six.iteritems(spec_dict) if is_bfloat_feature(v)
]
for c in to_convert:
result[c] = tf.cast(result[c], tf.bfloat16)
return result
def _generate_saved_model_for_half_plus_two(export_dir,
tf2=False,
as_text=False,
as_tflite=False,
as_tflite_with_sigdef=False,
use_main_op=False,
include_mlmd=False,
device_type="cpu"):
"""Generates SavedModel for half plus two.
Args:
export_dir: The directory to which the SavedModel should be written.
tf2: If True generates a SavedModel using native (non compat) TF2 APIs.
as_text: Writes the SavedModel protocol buffer in text format to disk.
as_tflite: Writes the Model in Tensorflow Lite format to disk.
as_tflite_with_sigdef: Writes the Model with SignatureDefs in Tensorflow
Lite format to disk.
use_main_op: Whether to supply a main op during SavedModel build time.
include_mlmd: Whether to include an MLMD key in the SavedModel.
device_type: Device to force ops to run on.
"""
if tf2:
hp = HalfPlusTwoModel()
tf.saved_model.save(hp, export_dir, signatures=hp.get_serving_signatures())
return
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
device_name = "/cpu:0"
if device_type == "gpu":
device_name = "/gpu:0"
with tf.Session(
graph=tf.Graph(),
config=tf.ConfigProto(log_device_placement=True)) as sess:
with tf.device(device_name):
# Set up the model parameters as variables to exercise variable loading
# functionality upon restore.
a = tf.Variable(0.5, name="a")
b = tf.Variable(2.0, name="b")
c = tf.Variable(3.0, name="c")
# Create a placeholder for serialized tensorflow.Example messages to be
# fed.
serialized_tf_example = tf.placeholder(
tf.string, name="tf_example", shape=[None])
# parse_example only works on CPU
with tf.device("/cpu:0"):
tf_example = tf.parse_example(serialized_tf_example,
_get_feature_spec())
if as_tflite:
# TFLite v1 converter does not support unknown shape.
x = tf.ensure_shape(tf_example["x"], (1, 1), name="x")
else:
# Use tf.identity() to assign name
x = tf.identity(tf_example["x"], name="x")
if as_tflite_with_sigdef:
# Resulting TFLite model will have input named "tflite_input".
x = tf.ensure_shape(tf_example["x"], (1, 1), name="tflite_input")
if device_type == "mkl":
# Create a small convolution op to trigger MKL
# The op will return 0s so this won't affect the
# resulting calculation.
o1 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
y = o1[0, 0, 0, 0] + tf.add(tf.multiply(a, x), b)
else:
y = tf.add(tf.multiply(a, x), b)
y = tf.identity(y, name="y")
if device_type == "mkl":
# Create a small convolution op to trigger MKL
# The op will return 0s so this won't affect the
# resulting calculation.
o2 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
y2 = o2[0, 0, 0, 0] + tf.add(tf.multiply(a, x), c)
else:
y2 = tf.add(tf.multiply(a, x), c)
y2 = tf.identity(y2, name="y2")
x2 = tf.identity(tf_example["x2"], name="x2")
if device_type == "mkl":
# Create a small convolution op to trigger MKL
# The op will return 0s so this won't affect the
# resulting calculation.
o3 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
y3 = o3[0, 0, 0, 0] + tf.add(tf.multiply(a, x2), c)
else:
# Add separate constants for x2, to prevent optimizers like TF-TRT from
# fusing the paths to compute y/y2 and y3 together.
a2 = tf.Variable(0.5, name="a2")
c2 = tf.Variable(3.0, name="c2")
y3 = tf.add(tf.multiply(a2, x2), c2)
y3 = tf.identity(y3, name="y3")
assign_filename_op = _create_asset_file()
predict_signature_def = _build_predict_signature(x, y)
signature_def_map = {
"regress_x_to_y":
_build_regression_signature(serialized_tf_example, y),
"regress_x_to_y2":
_build_regression_signature(serialized_tf_example, y2),
"regress_x2_to_y3":
_build_regression_signature(x2, y3),
"classify_x_to_y":
_build_classification_signature(serialized_tf_example, y),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
_build_predict_signature(x, y)
}
# Initialize all variables and then save the SavedModel.
sess.run(tf.global_variables_initializer())
if as_tflite or as_tflite_with_sigdef:
converter = tf.lite.TFLiteConverter.from_session(sess, [x], [y])
tflite_model = converter.convert()
if as_tflite_with_sigdef:
k = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
tflite_model = signature_def_utils.set_signature_defs(
tflite_model, {k: predict_signature_def})
open(export_dir + "/model.tflite", "wb").write(tflite_model)
else:
if use_main_op:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
main_op=tf.group(tf.saved_model.main_op.main_op(),
assign_filename_op))
else:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
main_op=tf.group(assign_filename_op))
if not as_tflite:
builder.save(as_text)
if include_mlmd:
_write_mlmd(export_dir, "test_mlmd_uuid")
