Exemplo n.º 1
0
    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
Exemplo n.º 2
0
 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(())
Exemplo n.º 3
0
 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(())
Exemplo n.º 4
0
 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))
Exemplo n.º 5
0
    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))