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, }