def test_log_prob(self):
d = tfd.Masked(tfd.Normal(tf.zeros([20]), 1), tf.sequence_mask(15, 20))
x = np.linspace(-.1, .1, 20).astype(np.float32)
self.assertAllClose(
tf.pad(tfd.Normal(0, 1).log_prob(x[:15]), [[0, 5]]),
d.log_prob(x))
d = tfd.Masked(tfd.Normal(tf.zeros([20]), 1), (tf.range(20) % 2) > 0)
x = np.linspace(-.1, .1, 20).astype(np.float32)
self.assertAllClose(tf.zeros(10), d.log_prob(x)[::2])
self.assertAllClose(tfd.Normal(0, 1).log_prob(x)[1::2], d.log_prob(x)[1::2])
def get_inputs_from_targets(self, targets, start_token_ids):
"""Converts target ids to input ids, i.e. adds <s> and removes last."""
length = tf.math.count_nonzero(targets, axis=1, dtype=tf.int32)
# Add start token ids.
inputs = tf.concat([tf.expand_dims(start_token_ids, axis=1), targets],
1)
# Remove </s> from the input.
mask = tf.sequence_mask(length,
self.params["max_decoder_length"] + 1,
dtype=tf.int32)
inputs = (mask * inputs)[:, :-1]
return inputs
def _to_nested_map(self, text) -> py_utils.NestedMap:
ids, labels, paddings = self.tokenizer.StringsToIds(
text, max_length=self.params.max_sequence_length)
# Unfortunately some tokenizers don't return the correct paddings.
# We recompute it by looking at when the labels sequence terminates.
indices = tf.where(tf.math.equal(labels, self.tokenizer.eos_id))
lengths = tf.math.segment_min(indices[:, 1], indices[:, 0]) + 1
new_paddings = tf.cast(
1.0 - tf.sequence_mask(
lengths,
maxlen=self.params.max_sequence_length,
dtype=paddings.dtype),
dtype=paddings.dtype)
return py_utils.NestedMap(ids=ids, labels=labels, paddings=new_paddings)
def call(self, inputs, lengths, training):
seq_mask = tf.sequence_mask(
lengths, inputs.shape[1], dtype=tf.dtypes.float32)
forward_outputs = self.forward_rnn(inputs, training=training)
reversed_inputs = tf.reverse_sequence(inputs, lengths, seq_axis=1)
backward_outputs = self.backward_rnn(reversed_inputs, training=training)
backward_outputs = tf.reverse_sequence(
backward_outputs, lengths, seq_axis=1)
outputs = tf.concat([forward_outputs, backward_outputs], axis=-1)
outputs = outputs * tf.expand_dims(seq_mask, -1)
if self.reduce_states:
outputs = tf.reduce_mean(outputs, axis=1)
return outputs
def call(self, inputs):
"""Generates mask (whether example is valid) from features.
Args:
inputs: (dict) Features with a mix of context (2D) and example features
(3D).
Returns:
mask: (tf.Tensor) Mask is a tensor of shape [batch_size, list_size], which
is True for a valid example and False for invalid one.
"""
example_feature = inputs[next(six.iterkeys(self._example_feature_columns))]
list_size = tf.shape(example_feature)[1]
sizes = inputs[self._size_feature_name]
mask = tf.sequence_mask(sizes, maxlen=list_size)
return mask
def testSaveTensorKwarg(self):
class LayerWithTensorKwarg(keras.layers.Layer):
def call(self, inputs, tensor=None):
if tensor is not None:
return inputs * tf.cast(tensor, tf.float32)
else:
return inputs
t = self.evaluate(tf.sequence_mask(1))
inputs = keras.layers.Input(shape=(3))
model = keras.models.Model(inputs, LayerWithTensorKwarg()(inputs, t))
input_arr = np.random.random((1, 3))
predictions = model.predict(input_arr)
saved_model_dir = self._save_model_dir()
model.save(saved_model_dir, save_format='tf')
loaded = keras_load.load(saved_model_dir)
loaded_predictions = loaded.predict(input_arr)
self.assertAllClose(predictions, loaded_predictions)
def test_functional_cloning_with_tensor_kwarg(self, share_weights):
"""Test that cloning works with models that use Tensor kwargs."""
if share_weights:
clone_fn = functools.partial(
keras.models.clone_model, clone_function=models.share_weights
)
else:
clone_fn = keras.models.clone_model
class LayerWithTensorKwarg(keras.layers.Layer):
def call(self, inputs, tensor=None):
if tensor is not None:
return inputs * tf.cast(tensor, tf.float32)
else:
return inputs
inputs = keras.layers.Input(shape=(3))
t = tf.sequence_mask(tf.shape(inputs)[1])
model = keras.models.Model(inputs, LayerWithTensorKwarg()(inputs, t))
model.add_loss(tf.reduce_sum(model.outputs))
input_arr = np.random.random((1, 3)).astype(np.float32)
clone = clone_fn(model)
if tf.executing_eagerly():
clone(input_arr)
loss = clone.losses[0]
else:
with self.session() as sess:
clone(input_arr)
if share_weights:
self.skipTest(
"Weight sharing with inputs in call **kwargs does "
"not work correctly in v1"
)
else:
feed_dict = {clone.input: input_arr}
loss = sess.run(clone.losses[0], feed_dict=feed_dict)
self.assertAllClose(np.sum(input_arr), loss)
def test_Bidirectional_with_time_major_input(self, time_major):
batch_size, time, input_dim = 2, 3, 1
inputs = tf.zeros((batch_size, time, input_dim))
# length is [1 2]. Within the batch, the first element has 1 step, and the
# second element as 2 steps.
lengths = tf.range(1, 1 + batch_size)
mask = tf.sequence_mask(lengths, maxlen=time, dtype=tf.float32)
forward_cell = _AddOneCell(name="forward")
backward_cell = _AddOneCell(name="backward")
layer = keras.layers.Bidirectional(
layer=keras.layers.RNN(forward_cell,
time_major=time_major,
return_sequences=True),
backward_layer=keras.layers.RNN(
backward_cell,
time_major=time_major,
return_sequences=True,
go_backwards=True,
),
)
# Switch to time-major.
if time_major:
inputs = tf.transpose(inputs, [1, 0, 2])
mask = tf.transpose(mask, [1, 0])
keras_outputs = layer(inputs, mask=mask)
if time_major:
keras_outputs = tf.transpose(keras_outputs, [1, 0, 2])
# expect the first element in batch has 1 step and second element in batch
# has 2 steps.
expected_result = np.array([
[[1.0, 1.0], [0.0, 0.0], [0.0, 0.0]],
[[1.0, 1.0], [1.0, 1.0], [0.0, 0.0]],
])
self.assertAllClose(expected_result, keras_outputs)
def test_compare_ragged_with_masks(self, layer):
vocab_size = 100
timestep = 20
units = 32
embedder = embeddings.Embedding(input_dim=vocab_size, output_dim=units)
layer = layer(units, return_sequences=True)
data = tf.constant(
np.random.RandomState(0).randint(0, vocab_size, [timestep, timestep]))
mask = tf.sequence_mask(tf.range(1, timestep + 1))
data_ragged = tf.ragged.boolean_mask(data, mask)
outputs = []
devices = [testing_utils.device(should_use_gpu=False)]
if tf.test.is_gpu_available():
devices.append(testing_utils.device(should_use_gpu=True))
for device in devices:
with device:
outputs.append(tf.boolean_mask(layer(embedder(data), mask=mask), mask))
outputs.append(layer(embedder(data_ragged)).values)
for i in range(len(outputs) - 1):
self.assertAllClose(outputs[i], outputs[i + 1], atol=1e-4)
def test_shape(self): d = tfd.Masked(tfd.Normal(tf.zeros([7, 20]), 1), tf.sequence_mask(15, 20)) self.assertEqual((), d.event_shape) self.assertAllEqual([], d.event_shape_tensor()) self.assertEqual((7, 20), d.batch_shape) self.assertAllEqual([7, 20], d.batch_shape_tensor())
def _parse_fn(record):
"""Parses a record into a feature_dict."""
feature_values = tf.io.parse_single_example(
serialized=record,
features={
'i/o':
tf.io.FixedLenFeature([], tf.string, default_value=''),
'program_encoding':
tf.io.FixedLenFeature([], tf.string, default_value=''),
})
ios = tf.strings.split(tf.strings.split(feature_values['i/o'],
sep='>'),
sep='<')
inputs, outputs = ios.merge_dims(0, 1)[::2], ios.merge_dims(0, 1)[1::2]
# Parse inputs into tokens.
inputs = tf.strings.unicode_split(inputs, 'UTF-8').to_tensor()
inputs = spec_vocab_table.lookup(inputs) # Map characters to tokens.
# Parse outputs into tokens.
outputs_with_separators = (tf.strings.unicode_split(
outputs, 'UTF-8').to_tensor())
outputs_with_separators = spec_vocab_table.lookup(
outputs_with_separators)
split_outputs = tf.strings.unicode_split(
tf.strings.split(outputs, sep='|'), 'UTF-8')
outputs = split_outputs.merge_dims(1, 2).to_tensor()
outputs = spec_vocab_table.lookup(outputs)
# Compute indices for the start of each part of the spec, w.r.t. the
# original spec.
separator_indices = tf.where(
tf.equal(outputs_with_separators, separator_id))[:, 1]
separator_indices = tf.reshape(
separator_indices, (tf.shape(outputs_with_separators)[0], -1))
start_indices = separator_indices - tf.expand_dims(
tf.range(tf.shape(separator_indices)[1], dtype=tf.int64), 0)
start_indices = tf.concat((tf.zeros(
(tf.shape(start_indices)[0], 1), dtype=tf.int64), start_indices),
axis=1)
num_examples = tf.shape(start_indices)[0]
num_parts = tf.shape(start_indices)[1]
# Construct the shifted spec suffixes.
flat_start_indices = tf.reshape(start_indices, (-1, ))
prefix_mask = (1 - tf.sequence_mask(
flat_start_indices, maxlen=tf.shape(outputs)[-1], dtype=tf.int64))
masked_outputs = tf.repeat(outputs, num_parts, axis=0) * prefix_mask
output_suffixes = tf.vectorized_map(
fn=lambda x: tf.roll(x[0], x[1], axis=0),
elems=(masked_outputs, -flat_start_indices))
# Compute indices for the start/end of spec parts, w.r.t. the shifted spec
# suffixes.
ground_truth_start_indices = tf.zeros((num_examples * num_parts, ),
dtype=tf.int64)
cumulative_end_indices = tf.concat(
(start_indices,
tf.math.count_nonzero(outputs, axis=-1, keepdims=True)),
axis=1)
ground_truth_end_indices = tf.reshape(
cumulative_end_indices[:, 1:] - cumulative_end_indices[:, :-1],
(-1, ))
# Construct the actual spec parts to predict.
range_indices = tf.expand_dims(tf.range(tf.shape(output_suffixes)[-1],
dtype=tf.int64),
axis=0)
part_mask = tf.where(
tf.logical_and(
range_indices >= tf.expand_dims(ground_truth_start_indices,
axis=1),
range_indices < tf.expand_dims(ground_truth_end_indices,
axis=1)), 1, 0)
output_parts = output_suffixes * tf.cast(part_mask, tf.int64)
output_parts = tf.pad(output_parts,
[[0, 0], [0, 1]]) # Make room for sep.
# TODO(kshi): roll output_parts leftward by start_indices for SCAN.
first_zero_index = tf.math.count_nonzero(output_parts, axis=-1)
output_parts += tf.one_hot(first_zero_index,
depth=tf.shape(output_parts)[-1],
dtype=tf.int64) * separator_id
# Reshape everything so that different spec suffixes become different
# dataset elements.
output_suffixes_reshaped = tf.transpose(
tf.reshape(output_suffixes, (num_examples, num_parts, -1)),
(1, 0, 2))
output_parts_reshaped = tf.transpose(
tf.reshape(output_parts, (num_examples, num_parts, -1)), (1, 0, 2))
inputs_reshaped = tf.reshape(tf.tile(inputs, (num_parts, 1)),
(num_parts, num_examples, -1))
ground_truth_start_indices_reshaped = tf.transpose(
tf.reshape(ground_truth_start_indices, (num_examples, num_parts)))
ground_truth_end_indices_reshaped = tf.transpose(
tf.reshape(ground_truth_end_indices, (num_examples, num_parts)))
# Combine spec parts from all examples into one sequence with separator
# tokens between examples and ending in EOS.
shifts = tf.cumsum(tf.concat((tf.zeros(
(num_parts, 1),
dtype=tf.int64), ground_truth_end_indices_reshaped[:, :-1] + 1),
1),
axis=-1)
flat_shifts = tf.reshape(shifts, (-1, ))
output_len = tf.shape(output_parts_reshaped)[-1]
flat_spec_parts = tf.reshape(output_parts_reshaped, (-1, output_len))
flat_spec_parts = tf.pad(flat_spec_parts,
[[0, 0], [0, max_target_length - output_len]])
combined_spec_parts = tf.vectorized_map(
fn=lambda x: tf.roll(x[0], x[1], axis=0),
elems=(flat_spec_parts, flat_shifts))
combined_spec_parts = tf.reshape(combined_spec_parts,
(num_parts, num_examples, -1))
combined_spec_parts = tf.reduce_sum(combined_spec_parts, axis=1)
first_zero_index = tf.math.count_nonzero(combined_spec_parts, axis=-1)
combined_spec_parts += tf.one_hot(
first_zero_index,
depth=tf.shape(combined_spec_parts)[-1],
dtype=tf.int64) * eos_id
# Create a dataset containing data for all spec suffixes.
dataset = tf.data.Dataset.from_tensor_slices({
'inputs':
inputs_reshaped,
'outputs':
output_suffixes_reshaped,
'spec_parts':
combined_spec_parts,
'start_index':
ground_truth_start_indices_reshaped,
'end_index':
ground_truth_end_indices_reshaped
})
return dataset
def decode_fn(value, data_aug=False,
max_num_points=245760, max_num_bboxes=100,
class_id=1, difficulty=1, pillar_map_size=(256, 256),
pillar_map_range=(75.2, 75.2)):
"""Decode function."""
tensor_dict = waymo_decoder.decode_tf_example(
serialized_example=value,
features=waymo_decoder.FEATURE_SPEC)
frame_name = tensor_dict['frame_name']
points_xyz = tensor_dict['lidars']['points_xyz']
points_feature = tensor_dict['lidars']['points_feature']
bboxes = tensor_dict['objects']['box']
bboxes_label = tensor_dict['objects']['label']
bboxes_speed = tensor_dict['objects']['speed']
bboxes_difficulty = tensor_dict['objects']['combined_difficulty_level']
num_valid_points = tf_util.get_shape(points_xyz)[0]
points_xyz = tf_util.pad_or_trim_to(points_xyz, [max_num_points, 3])
points_feature = tf_util.pad_or_trim_to(points_feature, [max_num_points, 2])
points_mask = tf.sequence_mask(num_valid_points,
maxlen=max_num_points)
points_mask = tf.cast(points_mask, dtype=tf.dtypes.float32)
bboxes_difficulty = bboxes_difficulty <= difficulty
bboxes_mask = tf.equal(bboxes_label, class_id)
bboxes_mask = tf.math.logical_and(bboxes_mask, bboxes_difficulty)
bboxes_mask = tf.cast(bboxes_mask, dtype=tf.dtypes.float32)
num_valid_bboxes = tf_util.get_shape(bboxes)[0]
bboxes_index = tf.math.top_k(
bboxes_mask, k=tf.math.minimum(max_num_bboxes, num_valid_bboxes))[1]
bboxes_mask = tf.gather(bboxes_mask, bboxes_index)
bboxes_label = tf.gather(bboxes_label, bboxes_index)
bboxes = tf.gather(bboxes, bboxes_index)
bboxes_speed = tf.gather(bboxes_speed, bboxes_index)
bboxes = tf_util.pad_or_trim_to(bboxes, [max_num_bboxes, 7])
bboxes_label = tf_util.pad_or_trim_to(bboxes_label, [max_num_bboxes])
bboxes_speed = tf_util.pad_or_trim_to(bboxes_speed, [max_num_bboxes, 2])
bboxes_difficulty = tf_util.pad_or_trim_to(bboxes_difficulty,
[max_num_bboxes])
bboxes_mask = tf_util.pad_or_trim_to(bboxes_mask, [max_num_bboxes])
if data_aug:
points_xyz, points_mask, bboxes = augment(
points_xyz=points_xyz,
points_mask=points_mask,
bboxes=bboxes)
(pillar_map_xyz, pillar_map_bboxes, pillar_map_bboxes_label,
pillar_map_if_in_bboxes, pillar_map_centerness, pillar_map_bboxes_index) = (
assign_bboxes(
pillar_map_size=pillar_map_size,
pillar_map_range=pillar_map_range,
bboxes=bboxes,
bboxes_label=bboxes_label,
bboxes_mask=bboxes_mask))
pillar_map_xyz = tf.reshape(pillar_map_xyz, [-1, 3])
pillar_map_bboxes = tf.reshape(pillar_map_bboxes, [-1, 7])
pillar_map_bboxes_label = tf.reshape(pillar_map_bboxes_label, [-1])
pillar_map_if_in_bboxes = tf.reshape(pillar_map_if_in_bboxes, [-1])
pillar_map_centerness = tf.reshape(pillar_map_centerness, [-1])
pillar_map_bboxes_index = tf.reshape(pillar_map_bboxes_index, [-1])
return {
'frame_name': frame_name,
'points_xyz': points_xyz,
'points_feature': points_feature,
'points_mask': points_mask,
'bboxes': bboxes,
'bboxes_label': bboxes_label,
'bboxes_mask': bboxes_mask,
'bboxes_speed': bboxes_speed,
'pillar_map_xyz': pillar_map_xyz,
'pillar_map_bboxes': pillar_map_bboxes,
'pillar_map_if_in_bboxes': pillar_map_if_in_bboxes,
}
