def offset_batched_neighbors(self): """ [B, n?, k?] offset index into flattened input batched cloud features. For point cloud features of size [B, n?, ...], `tf.gather(features.values, offset_batched_neighbors)` will give the same result as ``` tf.map(lambda features, batched_neighbors: tf.gather(features, batched_neighbors)) ``` The returned value is a model input. """ if self._offset_batched_neighbors is None: row_length = utils.leading_dim(self.in_coords) row_lengths = b.batched(row_length) batched_neighbors = b.batched(self.neighbors) # row_lengths = b.as_batched_model_input(row_length) # batched_neighbors = b.as_batched_model_input(self.neighbors) offset = tf.keras.layers.Lambda( tf.math.cumsum, arguments=dict(exclusive=True))(row_lengths) offset_batched_neighbors = utils.apply_row_offset( batched_neighbors, offset) self._offset_batched_neighbors = b.as_model_input( offset_batched_neighbors) return self._offset_batched_neighbors
def test_batched_tensor(self): b = builder.MetaNetworkBuilder() with b: x = b.prebatch_input(shape=(4, 5), dtype=tf.float32) batched = builder.batched(x) self.assertEqual(batched.shape.as_list(), [None, 4, 5]) b.preprocessor(())
def test_batched_ragged(self): b = builder.MetaNetworkBuilder() with b: x = b.prebatch_input(shape=(None, 5), dtype=tf.float32) batched = builder.batched(x) self.assertTrue(isinstance(batched, tf.RaggedTensor)) self.assertEqual(batched.shape.as_list(), [None, None, 5]) b.preprocessor(())
def batch_labels_and_weights(self, labels, weights=None): from more_keras.meta_models import builder as b return tf.nest.map_structure( lambda tensor: None if tensor is None else b.batched(tensor), (labels, weights))
def _test_transposed_consistent(self): batch_size = 3 with tf.device('/cpu:0'): np_data = [ np.random.uniform(size=(100, 3)).astype(np.float32), np.random.uniform(size=(200, 3)).astype(np.float32), np.random.uniform(size=(50, 3)).astype(np.float32), ] def generator(): labels = np.array([0, 1, 2], dtype=np.int64) indices = [ np.array([0, 5, 2, 7, 10], dtype=np.int64), np.array([1, 4, 3, 2], dtype=np.int64), np.array([10, 15], dtype=np.int64), ] yield (np_data[0], indices[0]), labels[0] yield (np_data[1], indices[1]), labels[1] yield (np_data[2], indices[2]), labels[2] dataset = tf.data.Dataset.from_generator( generator, output_types=((tf.float32, tf.int64), tf.int64), output_shapes=((tf.TensorShape( (None, 3)), tf.TensorShape((None,))), tf.TensorShape(()))) coords = b.prebatch_input((None, 3), tf.float32) sample_indices = b.prebatch_input((None,), dtype=tf.int64) labels = b.prebatch_input((), dtype=tf.int64) neighbors = q.query_pairs(coords, 0.1) in_place_neighborhood = n.InPlaceNeighborhood(coords, neighbors) sampled_neighborhood = n.SampledNeighborhood( in_place_neighborhood, sample_indices) transposed = sampled_neighborhood.transpose simple = n.Neighborhood(sampled_neighborhood.out_coords, sampled_neighborhood.in_coords, transposed.neighbors) simple_obn = simple.offset_batched_neighbors trans_obn = transposed.offset_batched_neighbors simple_out = b.as_batched_model_input(simple.in_coords) trans_out = b.as_batched_model_input(transposed.in_coords) out = [[o.flat_values, o.nested_row_splits] for o in (simple_obn, trans_obn, simple_out, trans_out)] flat_out = tf.nest.flatten(out) model = b.model(flat_out) preprocessor = b.preprocessor(b.batched(labels)) dataset = preprocessor.map_and_batch(dataset, batch_size) # if tf.executing_eagerly(): # for data, label in dataset.take(1): # pass # else: data, label = tf.compat.v1.data.make_one_shot_iterator( dataset).get_next() flat_out = model(tf.nest.flatten(data)) if not isinstance(flat_out, (list, tuple)): flat_out = flat_out, out = tf.nest.pack_sequence_as(out, flat_out) out = [tf.RaggedTensor.from_nested_row_splits(*o) for o in out] out, label = self.evaluate((out, label)) simple_obn, trans_obn, simple_coords, trans_coords = out self.assertRaggedEqual(simple_obn, trans_obn) self.assertRaggedEqual(simple_coords, trans_coords) self.assertEqual(simple_obn.nested_row_splits[-1][-1], sum(d.shape[0] for d in np_data))