def pad_or_truncate_pair(token_ids_a, token_ids_b, sequence_length, cls_id,
sep_id):
"""Pad or truncate pair."""
token_ids_a = token_ids_a[:sequence_length - 3]
truncated_len_a = tf.size(token_ids_a)
maximum_len_b = tf.maximum(sequence_length - 3 - truncated_len_a, 0)
token_ids_b = token_ids_b[:maximum_len_b]
truncated_len_b = tf.size(token_ids_b)
truncated_len_pair = truncated_len_a + truncated_len_b
padding = tf.zeros([sequence_length - 3 - truncated_len_pair], tf.int32)
token_ids = tf.concat([
[cls_id],
token_ids_a,
[sep_id],
token_ids_b,
[sep_id],
padding], 0)
mask = tf.concat([
tf.ones([truncated_len_pair + 3], tf.int32),
padding], 0)
segment_ids = tf.concat([
tf.zeros([truncated_len_a + 2], tf.int32),
tf.ones([truncated_len_b + 1], tf.int32),
padding], 0)
token_ids = tf.ensure_shape(token_ids, [sequence_length])
mask = tf.ensure_shape(mask, [sequence_length])
segment_ids = tf.ensure_shape(segment_ids, [sequence_length])
return token_ids, mask, segment_ids
def pad_or_truncate(token_ids, sequence_length, cls_id, sep_id): token_ids = token_ids[:sequence_length - 2] truncated_len = tf.size(token_ids) padding = tf.zeros([sequence_length - 2 - truncated_len], tf.int32) token_ids = tf.concat([[cls_id], token_ids, [sep_id], padding], 0) mask = tf.concat([tf.ones([truncated_len + 2], tf.int32), padding], 0) token_ids = tf.ensure_shape(token_ids, [sequence_length]) mask = tf.ensure_shape(mask, [sequence_length]) return token_ids, mask
def build_planner_inputs(question, answer, length, lookup_table):
"""Convert text to TextInputs for conditional text planner.
Args:
question: <string>, space-separated token string.
answer: <string>, space-separated token string.
length: Length to pad or truncate to.
lookup_table: Instance of contrib.lookup.index_table_from_tensor.
Returns:
Instance of TextInputs.
"""
# Build question.
q_tokens = tf.string_split([question]).values
q_tokens = tf.concat([["[Q]"], q_tokens], axis=0)
q_token_ids = tf.cast(lookup_table.lookup(q_tokens), tf.int32)
q_len = tensor_utils.shape(q_token_ids, 0)
q_positions = tf.range(q_len)
# Build answer.
a_tokens = tf.string_split([answer]).values
a_tokens = tf.concat([["[A]"], a_tokens], axis=0)
a_token_ids = tf.cast(lookup_table.lookup(a_tokens), tf.int32)
a_len = tensor_utils.shape(a_token_ids, 0)
a_positions = tf.range(a_len)
# Combine.
token_ids = tf.concat([q_token_ids, a_token_ids], axis=0)
segment_ids = tf.concat([tf.fill([q_len], 2), tf.fill([a_len], 1)], axis=0)
positions = tf.concat([q_positions, a_positions], axis=0)
q_mask = tf.ones_like(q_token_ids)
mask = tf.concat([q_mask, tf.ones_like(a_token_ids)], axis=0)
# Truncate.
token_ids = token_ids[:length]
segment_ids = segment_ids[:length]
mask = mask[:length]
positions = positions[:length]
# Pad.
pad = [[0, length - tf.size(token_ids)]]
token_ids = tf.pad(token_ids, pad)
mask = tf.pad(mask, pad)
segment_ids = tf.pad(segment_ids, pad)
positions = tf.pad(positions, pad)
text_input = TextInputs(token_ids=tf.ensure_shape(token_ids, [length]),
mask=tf.ensure_shape(mask, [length]),
segment_ids=tf.ensure_shape(segment_ids, [length]),
positions=tf.ensure_shape(positions, [length]))
return text_input
def pad_or_truncate_pair(token_ids, sequence_length):
"""Pad or truncate pair."""
token_ids = token_ids[:sequence_length]
truncated_len = tf.size(token_ids)
padding = tf.zeros([sequence_length - truncated_len], tf.int32)
token_ids = tf.concat([token_ids, padding], 0)
mask = tf.concat([tf.ones([truncated_len], tf.int32), padding], 0)
if FLAGS.encode_query:
segment_ids = tf.concat([tf.zeros([truncated_len], tf.int32), padding], 0)
else:
segment_ids = tf.concat([tf.ones([truncated_len], tf.int32), padding], 0)
token_ids = tf.ensure_shape(token_ids, [sequence_length])
mask = tf.ensure_shape(mask, [sequence_length])
segment_ids = tf.ensure_shape(segment_ids, [sequence_length])
return token_ids, mask, segment_ids
def decode_and_resize_img(img_bytes, height, width, depth):
# type: (tf.Tensor, int, int, int) -> tf.Tensor
"""Decodes and resizes input image to return a float representation.
Args:
img_bytes: tensor for input bytes.
height: Desired height of image.
width: Desired width of image.
depth: Desired bit depth of image.
Returns:
float_pixels: Tensor storing the image as float. This is the input tensor
that we'll reference in the Explainable AI feature to show how output
changes with input.
"""
features = tf.squeeze(img_bytes, axis=1, name='input_squeeze')
float_pixels = tf.map_fn(
# pylint: disable=g-long-lambda
lambda img: tf.image.resize_with_crop_or_pad(
tf.io.decode_image(img, channels=depth, dtype=tf.float32),
height, width),
features,
dtype=tf.float32,
name='input_convert')
float_pixels = tf.ensure_shape(float_pixels,
(None, height, width, depth))
float_pixels = tf.identity(float_pixels, name=input_pixels_tensor_name)
return float_pixels
def build_text_inputs(
text,
length,
lookup_table,
segment_id=0,
start_token=None,
end_token=None,
):
"""Convert text to TextInputs.
Args:
text: <string>, space-separated token string.
length: Length to pad or truncate to.
lookup_table: Instance of contrib.lookup.index_table_from_tensor.
segment_id: Integer denoting segment type.
start_token: Optional start token.
end_token: Optional end token.
Returns:
Instance of TextInputs.
"""
# Tokenize and truncate.
tokens = tf.string_split([text]).values
length_offset = sum(
[0 if i is None else 1 for i in [start_token, end_token]])
tokens = tokens[:length - length_offset]
if start_token is not None:
tokens = tf.concat([[start_token], tokens], axis=0)
if end_token is not None:
tokens = tf.concat([tokens, [end_token]], axis=0)
token_ids = tf.cast(lookup_table.lookup(tokens), tf.int32)
mask = tf.ones_like(token_ids)
segment_ids = tf.fill(tf.shape(token_ids), segment_id)
pad = [[0, length - tf.size(token_ids)]]
token_ids = tf.pad(token_ids, pad)
mask = tf.pad(mask, pad)
segment_ids = tf.pad(segment_ids, pad)
positions = tf.range(length)
text_input = TextInputs(token_ids=tf.ensure_shape(token_ids, [length]),
mask=tf.ensure_shape(mask, [length]),
segment_ids=tf.ensure_shape(segment_ids, [length]),
positions=tf.ensure_shape(positions, [length]))
return text_input
def parse_object_features(features, positions, params):
"""Parse ObjectDetectionOutput from TensorProtos."""
features = tf.io.parse_tensor(features, tf.float32)
positions = tf.io.parse_tensor(positions, tf.int64)
positions = tf.cast(positions, tf.int32)
features = features[:params["num_image_regions"]]
num_objects = tensor_utils.shape(features, 0)
padding = tf.maximum(0, params["num_image_regions"] - num_objects)
features = tf.pad(features, [[0, padding], [0, 0]])
positions = tf.pad(positions, [[0, padding]])
features = tf.ensure_shape(
features, [params["num_image_regions"], params["image_feature_size"]])
positions = tf.ensure_shape(positions, [params["num_image_regions"]])
mask = tf.pad(tf.ones(num_objects, dtype=tf.int32), [[0, padding]])
mask = tf.ensure_shape(mask, [params["num_image_regions"]])
output = ObjectDetectionOutput(features=features,
positions=positions,
mask=mask)
return output
def call(self, grid, pts, training=False):
"""Forward method for Learnable Voxel Grid.
Args:
grid: `[batch_size, *self.size, in_features]` tensor, input feature grid.
pts: `[batch_size, num_points, dim]` tensor, coordinates of points that
are within the range (0, 1).
training: bool, flag indicating training phase.
Returns:
outputs: `[batch_size, num_points, out_features]` tensor, continuous
function field value at locations specified at pts.
Raises:
RuntimeError: dimensions of grid does not match that of self.
"""
# assert that dimensions match
grid = tf.ensure_shape(
grid, (None, self.size[0], self.size[1], self.size[2], self.cin))
pts = tf.ensure_shape(pts, (None, None, self.dim))
lat, weights, xloc = self._interp(grid, pts)
outputs = self._eval_net(lat, weights, xloc, training=training)
return outputs
def parse_examples(serialized_example):
"""Make retrieval examples."""
feature_spec = dict(input_ids=tf.FixedLenSequenceFeature([], tf.int64,
True),
key=tf.FixedLenSequenceFeature([], tf.int64, True))
features = tf.parse_single_example(serialized_example, feature_spec)
features = {k: tf.cast(v, tf.int32) for k, v in features.items()}
block_ids, block_mask, block_segment_ids = pad_or_truncate_pair(
token_ids=features["input_ids"], sequence_length=FLAGS.block_seq_len)
key = tf.ensure_shape(features["key"], [1])
return dict(block_ids=block_ids,
block_mask=block_mask,
block_segment_ids=block_segment_ids,
key=key)
def recv(self, name, dtype=tf.float32, require_grad=False, shape=None):
with tf.control_dependencies([self._example_ids]):
tensor = self._bridge.receive_op(name, dtype)
if shape:
tensor = tf.ensure_shape(tensor, shape)
else:
logging.warning(
'Receiving tensor %s without checking shape. '
'Consider setting shape at model.recv(shape=(...)). '
'shape can have None dimensions '
'which matches to any length.', name)
self._train_ops.append(tensor)
self._recvs.append((name, tensor, require_grad))
return tensor
def preprocess_mapper(features, params, lookup_table, vocab):
"""Model-specific preprocessing of features from the dataset."""
# Set input type.
features["input_type"] = tf.constant(datasets.DatasetTypes.VQA)
# Fix question id.
features["question_id"] = tf.ensure_shape(features["question_id"], [])
features["question_inputs"] = text_utils.build_text_inputs(
text=features["question"],
length=params["question_length"],
lookup_table=lookup_table,
segment_id=0)
answer = text_utils.build_text_inputs(
text=features["answer"],
length=params["answer_length"],
lookup_table=lookup_table,
segment_id=1,
start_token=vocab.ANS)
assert isinstance(answer, text_utils.TextInputs)
features["answer_inputs"] = answer
features["answer_outputs"] = text_utils.TextOutputs(
token_ids=answer.token_ids[1:], mask=answer.mask[1:])
features["object_features"] = image_utils.parse_object_features(
features["object_features"], features["object_positions"], params)
if params.get("conditional_decoding"):
features["condition_inputs"] = text_utils.build_planner_inputs(
question=features["question"],
answer=features["answer"],
length=params["condition_length"],
lookup_table=lookup_table)
# Remove extra inputs.
features = {f: features[f] for f in features if f in KEYS}
# Add dummy inputs for standardization for multi-tasking.
for k, v in datasets.footprint(params).items():
if k not in features:
features[k] = v
return features
def _preprocess_image_for_eval(image,
preprocessing_options,
dataset_options,
bfloat16_supported=False):
"""Preprocesses an image for evaluation.
Args:
image: Either a rank 3 `Tensor` representing an image with shape [height,
width, channels=3] or else an encoded jpeg image string. The value of
`preprocessing_options.decode_input` should be set accordingly.
preprocessing_options: An instance of hparams.ImagePreprocessing.
dataset_options: An instance of DatasetOptions.
bfloat16_supported: Whether bfloat16 is supported on this platform.
Returns:
A preprocessed image `Tensor` of shape
[preprocessing_options.image_size, preprocessing_options.image_size, 3]
and dtype is tf.float32 or tf.bfloat16 depending on
`preprocessing_options.allow_mixed_precision` and `bfloat16_supported`. The
image values are in the range [-1, 1].
"""
with tf.name_scope('preprocess_for_eval'):
image_patch = _crop_to_square(
image,
decode_image=dataset_options.decode_input,
side_length=preprocessing_options.image_size,
crop_padding=preprocessing_options.crop_padding,
is_training=False,
resize_only=False,
eval_crop_method=preprocessing_options.eval_crop_method)
assert image_patch.dtype == tf.uint8, (
f'Expected uint8, was {image_patch.dtype}')
image_patch = _convert_image_dtype(
image_patch,
tf.bfloat16 if preprocessing_options.allow_mixed_precision
and bfloat16_supported else tf.float32)
return tf.ensure_shape(image_patch, [
preprocessing_options.image_size, preprocessing_options.image_size,
3
])
def _extract_summary(self, hidden_states, annotation_begins,
annotation_ends, annotation_labels, token_ids):
if self.mode == "cls":
return tf.expand_dims(hidden_states[:, 0, :], 1)
elif self.mode == "entity":
# [batch_size, local_num_summaries, 2 * hidden_size]
x = tf.concat([
tf.gather(hidden_states, annotation_begins, batch_dims=1),
tf.gather(hidden_states, annotation_ends, batch_dims=1)
],
axis=2)
# [batch_size, local_num_summaries, hidden_size]
x = self.extraction_linear(x)
# [batch_size, local_num_summaries, 1]
annotation_mask = tf.expand_dims(
tf.cast(tf.not_equal(annotation_labels, 0), tf.float32), -1)
annotation_mask = tf.ensure_shape(annotation_mask, [None, None, 1])
# [batch_size, local_num_summaries, hidden_size]
x = x * annotation_mask
return x
elif self.mode == "text_block":
token_block_offset = self.text_block_extract_every_x - 1
token_block_len = self.text_block_extract_every_x
x = tf.concat([
hidden_states[:, ::token_block_len, :],
hidden_states[:, token_block_offset::token_block_len, :]
],
axis=2)
x = self.extraction_linear(x)
labels = tf.cast(
tf.logical_and(
tf.not_equal(token_ids[:, ::token_block_len], 0),
tf.not_equal(
token_ids[:, token_block_offset::token_block_len], 0),
), tf.float32)
x = x * tf.expand_dims(labels, -1)
return x
else:
raise ValueError("Unknown summary mode: {}".format(self.mode))
def _validate_image_dimensions(image):
"""Verify that the image dimensions are valid.
Checks that the image has 3 color channels and is 3-dimensional. Updates
the static shape to have 3 channels, if the channel dimension was previously
unknown.
Args:
image: `Tensor` of raw image pixel data that is expected to be 3-dimensional
and have 3 channels in the last dimension. Can be any numeric type.
Returns:
The input image with the final dimension statically set to 3 if it was
previously None.
Raises:
tf.errors.InvalidArgumentError at runtime if the image is not 3-dimensional
or does not have 3 channels in the last dimension.
"""
with tf.name_scope('validate_image_dims'):
image = tf.ensure_shape(image, [None, None, 3])
return image
def _crop_to_square(image,
decode_image=False,
side_length=IMAGE_SIZE,
crop_padding=CROP_PADDING,
area_range=(0.08, 1.0),
is_training=True,
resize_only=False,
eval_crop_method=enums.EvalCropMethod.RESIZE_THEN_CROP):
"""Produces a (possibly distorted) square crop of an image.
Given an input image, either as an encoded bytes string or a decoded image
Tensor, produces a square version of it with the desired side length, using a
combination of cropping and resizing.
If `resize_only` is True, simply resize the image to be
`side_length`x`side_length`, possibly distorting it if the original image is
not square.
If `is_training` is True, then sample a random box to crop from the image and
then resize the result to be `side_length`x`side_length`.
If `is_training` is False then we follow `eval_crop_method` to determine the
strategy of cropping and resizing. Generally the approach is to end up with a
center crop of size `side_length`x`side_length` taken from the image resized
to have a minimum dimension of `side_length` + `crop_padding`. By setting
eval_crop_method appropriately, this can be accomplished by first resizing and
then cropping, first cropping and then resizing, or a less common approach of
cropping the central `side_length`/(`side_length`+`crop_padding`) pixels in
each dimension followed by resizing (and distorting) to
`side_length`x`side_length`.
If `decode_image` is True (i.e., `image` is an encoded jpeg image string),
when possible we crop before decoding, which can provide substantial speedups.
Args:
image: An image represented either as a 3D Tensor with any numeric DType or
else as an encoded jpeg image string.
decode_image: Whether `image` is an encoded jpeg image string or not.
side_length: The side length, in both spatial dimentions, of the output
image.
crop_padding: When `is_training` is False, this determines how much padding
to apply around the central square crop.
area_range: List of floats. The cropped area of the image must contain a
fraction of the supplied image within this range. Only relevant when
`is_training` is True and `resize_only` is False.
is_training: Whether this should operate in training (non-deterministic
random crop window) or eval (deterministic central crop window) mode.
resize_only: Whether to just resize the image to the target `side_length`
without performing any cropping. This is likely to distort the image.
eval_crop_method: The strategy for obtaining the desired square crop in eval
mode. See EvalCropMethod for valid values.
Returns:
An image Tensor of shape [`side_length`, `side_length`, 3]. If `image` was
provided then the output has the same dtype as `image`. If `image_bytes` was
provided then the output dtype is tf.uint8.
Raises:
ValueError: If both or neither of `image` and `image_bytes` was passed.
"""
with tf.name_scope('crop_to_square'):
if not decode_image:
image = _validate_image_dimensions(image)
if resize_only:
if decode_image:
image = _decode_and_maybe_crop_image(image)
resized = _resize_image(image, (side_length, side_length))
return tf.ensure_shape(resized, [side_length, side_length, 3])
image_shape = (tf.shape(image) if not decode_image else
tf.image.extract_jpeg_shape(image))
if is_training:
# During training, always crop then resize.
crop_window = _distorted_crop_window(image_shape,
area_range=area_range)
if decode_image:
cropped = _decode_and_maybe_crop_image(
image, _convert_3d_crop_window_to_2d(crop_window))
else:
cropped = tf.slice(image, crop_window[:3], crop_window[3:])
resized = _resize_image(cropped, [side_length, side_length])
return tf.ensure_shape(resized, [side_length, side_length, 3])
else:
# For eval, the ordering depends on eval_crop_method.
crop_frac = (side_length / (side_length + crop_padding))
if eval_crop_method == enums.EvalCropMethod.RESIZE_THEN_CROP:
if decode_image:
image = _decode_and_maybe_crop_image(image)
resize_dim = side_length + crop_padding
resized = _resize_to_min_dim(image, resize_dim)
crop_window = _center_crop_window(tf.shape(resized),
crop_dim=side_length)
cropped = tf.slice(resized, crop_window[:3], crop_window[3:])
return tf.ensure_shape(cropped, [side_length, side_length, 3])
elif eval_crop_method == enums.EvalCropMethod.CROP_THEN_RESIZE:
crop_window = _center_crop_window(image_shape,
crop_frac=crop_frac)
if decode_image:
cropped = _decode_and_maybe_crop_image(
image, _convert_3d_crop_window_to_2d(crop_window))
else:
cropped = tf.slice(image, crop_window[:3], crop_window[3:])
resized = _resize_image(cropped, [side_length, side_length])
return tf.ensure_shape(resized, [side_length, side_length, 3])
elif eval_crop_method == enums.EvalCropMethod.CROP_THEN_DISTORT:
if decode_image:
image = _decode_and_maybe_crop_image(image)
# Note that tf.image.central_crop does not produce a square crop. It
# preserves the input aspect ratio.
cropped = tf.image.central_crop(image,
central_fraction=crop_frac)
resized = _resize_image(cropped, [side_length, side_length])
return tf.ensure_shape(resized, [side_length, side_length, 3])
elif eval_crop_method == enums.EvalCropMethod.IDENTITY:
if decode_image:
image = _decode_and_maybe_crop_image(image)
return tf.ensure_shape(image, [side_length, side_length, 3])
def main(args):
graph = mtf.Graph()
mesh = mtf.Mesh(graph, "my_mesh")
batch_dim = mtf.Dimension("batch", args.batch_size)
row_blocks_dim = mtf.Dimension("row_blocks", 4)
col_blocks_dim = mtf.Dimension("col_blocks", 4)
rows_dim = mtf.Dimension("rows_size", 8)
cols_dim = mtf.Dimension("cols_size", 8)
# classes_dim = mtf.Dimension("classes", 10)
one_channel_dim = mtf.Dimension("one_channel", 3)
is_train = tf.placeholder(tf.bool, name="condition")
# ------------------ DATASET ----------------
def load(x):
x = np.load(x.decode())[np.newaxis, ...].astype(np.float32) / 1024 - 1
return x
data_path = os.path.join(args.dataset_root,
f'{args.image_size}x{args.image_size}/')
# retrieve dataseet for training
dataset_train, npy_data_train = npy_data(data_path,
args.scratch_path,
train_size=args.train_size,
train=True,
copy_files=True,
is_correct_phase=True)
dataset_train = dataset_train.map(
lambda x: tuple(tf.py_func(load, [x], [tf.float32])),
num_parallel_calls=AUTOTUNE)
dataset_train = dataset_train.batch(args.batch_size, drop_remainder=True)
dataset_train = dataset_train.repeat()
dataset_train = dataset_train.prefetch(AUTOTUNE)
dataset_train = dataset_train.make_one_shot_iterator()
# retrieve dataseet for testing
dataset_test, npy_data_test = npy_data(data_path,
args.scratch_path,
train_size=args.train_size,
train=False,
copy_files=True,
is_correct_phase=True)
dataset_test = dataset_test.map(
lambda y_test: tuple(tf.py_func(load, [y_test], [tf.float32])),
num_parallel_calls=AUTOTUNE)
dataset_test = dataset_test.batch(args.batch_size, drop_remainder=True)
dataset_test = dataset_test.repeat()
dataset_test = dataset_test.prefetch(AUTOTUNE)
dataset_test = dataset_test.make_one_shot_iterator()
# tensorflow if condition depending on is_train value. If True: dataset_train.get_next(). If False: dataset_test.get_next()
image_input = tf.cond(is_train, lambda: dataset_train.get_next(),
lambda: dataset_test.get_next())
image_input = tf.ensure_shape(image_input, [
args.batch_size, args.image_channels, args.image_size // 4,
args.image_size, args.image_size
])
image_input = mtf.import_tf_tensor(
mesh, tf.reshape(image_input, [args.batch_size, 1, 8, 32, 32]),
mtf.Shape([
batch_dim, row_blocks_dim, rows_dim, col_blocks_dim, cols_dim,
one_channel_dim
]))
# image_input = mtf.import_tf_tensor(
# mesh, tf.reshape(image_input, [args.batch_size, 4, 7, 4, 7, 1]),
# mtf.Shape([batch_dim, row_blocks_dim, rows_dim, col_blocks_dim, cols_dim, one_channel_dim]))
forward(image_input, args)
def _generate_saved_model_for_half_plus_two(export_dir,
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.
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.
"""
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)
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()
# 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 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")
def __init__(self, observation_size, net_arch, initializer, activation,
clip_range, value_coef, entropy_coef, learning_rate,
pre_training_learning_rate, action_bounds, policy):
"""
:param observation_size:
:param net_arch:
:param initializer:
:param activation:
:param clip_range:
:param value_coef:
:param entropy_coef:
:param learning_rate:
:param pre_training_learning_rate:
:param action_bounds:
:param policy:
"""
"""Set class constants"""
self.observation_size = observation_size
self.net_arch = net_arch
self.initializer = initializer
self.activation = activation
self.clip_range = clip_range
self.value_coef = value_coef
self.entropy_coef = entropy_coef
if action_bounds is None:
action_bounds = [0.0, 1.5]
self.action_bounds = action_bounds
self.learning_rate = learning_rate
self.pre_training_learning_rate = pre_training_learning_rate
if policy is None:
policy = GaussFull()
self.policy = policy
"""Set up the tensorflow graph"""
self.graph = Graph()
with self.graph.as_default():
self.sess = Session(graph=self.graph)
""" core """
# place holders
self.observation_string_ph = placeholder(
shape=(None, 1), dtype=string, name="observation_string_ph")
self.action_ph = placeholder(dtype=float32,
shape=(None, 1),
name="action_ph")
self.old_neg_logits = placeholder(dtype=float32,
shape=(None, 1),
name="old_neg_logits")
self.advantage_ph = placeholder(dtype=float32,
shape=(None, 1),
name="advantage_ph")
self.value_target_ph = placeholder(dtype=float32,
shape=(None, 1),
name="value_target_ph")
# learning rate tensors
self.learning_rate_ph = placeholder_with_default(
input=self.learning_rate, shape=())
self.pre_training_learning_rate_ph = placeholder_with_default(
input=self.pre_training_learning_rate, shape=())
# observation tensor
replaced1 = regex_replace(self.observation_string_ph, "/", "_")
replaced2 = regex_replace(replaced1, r"\+", "-")
byte_tensor = decode_base64(replaced2)
decoded = decode_raw(byte_tensor, out_type=float32)
squeezed = squeeze(decoded, axis=1)
self.observation_input = ensure_shape(
squeezed,
shape=(None, self.observation_size),
name="observation_input")
# policy net
latent_policy = net_core(self.observation_input, self.net_arch,
self.initializer, self.activation)
self.policy.construct(latent_policy=latent_policy)
self.clipped_action = clip_by_value(
cast(self.policy.action, float32), self.action_bounds[0],
self.action_bounds[1], "clipped_action")
# value net
latent_value = net_core(self.observation_input, self.net_arch,
self.initializer, self.activation)
self.value = identity(
input=Dense(units=1,
activation=None,
kernel_initializer=self.initializer)(latent_value),
name="value")
"""loss calculation"""
# policy loss
self.neg_logits = self.policy.neg_logits_from_actions(
self.action_ph)
ratio = exp(self.old_neg_logits - self.neg_logits)
standardized_adv = (self.advantage_ph - reduce_mean(
self.advantage_ph)) / (reduce_std(self.advantage_ph) + 1e-8)
raw_policy_loss = -standardized_adv * ratio
clipped_policy_loss = -standardized_adv * clip_by_value(
ratio, 1 - self.clip_range, 1 + self.clip_range)
self.policy_loss = reduce_mean(
maximum(raw_policy_loss, clipped_policy_loss))
self.value_loss = mean_squared_error(self.value_target_ph,
self.value)
# entropy loss
self.entropy_loss = -reduce_mean(self.policy.entropy)
# total loss
self.total_loss = self.policy_loss + self.value_coef * self.value_loss + self.entropy_coef * self.entropy_loss
# optimizer
optimizer = AdamOptimizer(learning_rate=self.learning_rate_ph)
# training ops
self.training_op = optimizer.minimize(self.total_loss)
# pre training
self.dist_param_target_ph = placeholder(
dtype=float32,
shape=(None, self.policy.dist_params.shape[1]),
name="dist_param_label_ph")
self.pre_training_loss = mean_squared_error(
self.dist_param_target_ph, self.policy.dist_params)
pre_training_optimizer = GradientDescentOptimizer(
learning_rate=self.pre_training_learning_rate_ph)
self.pre_training_op = pre_training_optimizer.minimize(
self.pre_training_loss)
"""utility nodes"""
# inspect model weights
self.trainable_variables = trainable_variables()
# saviour
self.saver = Saver()
# tensorboard summaries
self.summary = merge([
histogram("values", self.value),
histogram("advantages", standardized_adv),
histogram("actions", self.clipped_action),
histogram("det_actions",
replace_nan(self.policy.det_action, 0.0)),
histogram("value_targets", self.value_target_ph),
scalar("policy_loss", self.policy_loss),
scalar("value_loss", self.value_loss),
scalar("entropy_loss", self.entropy_loss)
])
self.pre_summary = merge([
histogram("pretraining_actions", self.clipped_action),
scalar("pretraining_loss", self.pre_training_loss)
])
# initialization
init = global_variables_initializer()
self.sess.run(init)
