def testNeighborSquaredDistance(self): n, p1, k = 2, 10, 3 points = tf.random.uniform((n, p1, 3)) neighbor_idx = tf.random.uniform((n, p1, k), minval=0, maxval=p1, dtype=tf.int32) neighbor_points = car_lib.MatmulGather(points, neighbor_idx) sq_dist_result = car_lib.NeighborSquaredDistanceMatrix( points, neighbor_points) with self.session() as sess: [np_points, np_neighbor_idx, np_sq_dist_result] = sess.run([points, neighbor_idx, sq_dist_result]) np_sq_dist_expected = self._np_sq_dis_neighbors(np_points, np_neighbor_idx) self.assertAllClose(np_sq_dist_result, np_sq_dist_expected)
def FProp(self, theta, input_data): """Apply projection to inputs. Args: theta: A NestedMap object containing weights' values of this layer and its children layers. input_data: A NestedMap object containing 'points', 'features', 'padding' Tensors, all of type tf.float32. 'points': Shape [N, P1, 3] 'features': Shape [N, P1, F] 'padding': Shape [N, P1] where 0 indicates real, 1 indicates padded. Returns: A NestedMap consisting of the following two NestedMaps, grouped_points: consists of the grouped points, features and padding. query_points: consists of the sampled points and padding. """ p = self.params features = input_data.features n, p1, c = py_utils.GetShape(features) points = py_utils.HasShape(input_data.points, [n, p1, 3]) padding = py_utils.HasShape(input_data.padding, [n, p1]) # Sampling sampled_idx, _ = car_lib.FarthestPointSampler( points, padding, num_sampled_points=p.num_samples) query_points = car_lib.MatmulGather(points, tf.expand_dims(sampled_idx, -1)) query_points = tf.squeeze(query_points, -2) # Grouping grouped_idx, grouped_padding = car_lib.NeighborhoodIndices( points, query_points, p.group_size, points_padding=padding, max_distance=p.ball_radius, sample_neighbors_uniformly=p.sample_neighbors_uniformly) grouped_points = car_lib.MatmulGather(points, grouped_idx) # Normalize the grouped points based on the location of the query point. grouped_points -= tf.expand_dims(query_points, -2) grouped_features = car_lib.MatmulGather(features, grouped_idx) # Get the padding for the query points. query_padding = tf.batch_gather(padding, sampled_idx) # Verify the shapes of output tensors. query_points = py_utils.HasShape(query_points, [n, p.num_samples, 3]) query_padding = py_utils.HasShape(query_padding, [n, p.num_samples]) grouped_features = py_utils.HasShape( grouped_features, [n, p.num_samples, p.group_size, c]) grouped_padding = py_utils.HasShape(grouped_padding, [n, p.num_samples, p.group_size]) output_grouped_points = py_utils.NestedMap(points=grouped_points, features=grouped_features, padding=grouped_padding) output_query = py_utils.NestedMap(points=query_points, padding=query_padding) output_map = py_utils.NestedMap({ 'grouped_points': output_grouped_points, 'query_points': output_query }) return output_map