lingvo.tasks.car.car_lib module

Library of functions on tensors for car layers, builders, and models.

lingvo.tasks.car.car_lib.SquaredDistanceMatrix(pa, pb, mem_optimized=False)

Compute pair-wise squared distances.

Expanded version (faster but potentially numerically unstable):

distance = pa^2 - 2*pa*pb + pb^2

Non-expanded version (slow but numerically stable):

distance = (pa - pb)^2

Parameters

• pa – tensor of shape [N, P1, dims]

• pb – tensor of shape [N, P2, dims]

• mem_optimized – Whether to use the memory-optimized expanded formulation. Defaults to False. If enabled, the expanded version is used that may have numerical issues.

Returns

tensor of shape [N, P1, P2]

lingvo.tasks.car.car_lib.NeighborSquaredDistanceMatrix(points, neighbor_points)

Compute the squared distance matrix between points and their neighbors.

Parameters

• points – A float tf.Tensor of shape [N, P1, 3] with point positions.

• neighbor_points – A float tf.Tensor fo shape [N, P1, K, 3] with neighbor positions.

Returns

Squared distance matrix between points and their K nearest neighbors as a float tf.Tensor of shape [N, P1, K].

lingvo.tasks.car.car_lib.KnnIndices(points, query_points, k, valid_num=None, max_distance=None)

k-nearest neighbors of query_points in points.

The caller should ensure that points[i, :valid_num[i], :] are the non-padding points.

Padding is returned alongside indices. Non-padded points are guaranteed to be unique (non-repeated) points from original non-padded points.

Padded points arise due to either a lack of points (k exceeds valid_num) or points are too far away (exceeds max distance).

TODO(weihan,jngiam): For backwards compatibility with PointCNN, if there are fewer than k points to select (possibly because of valid_num), the points selected will first consist of those in the non-padded points, and then those from the padded points. This assumes that the padded points are duplications of the original points. PointCNN should be updated to respect padding.

The auxiliary input ‘valid_num’ marks the number of non-padding points in each sample. This is needed because we randomly duplicated points to make the input fix-sized, we want search for k-NN in non-padding points first otherwise the result may degenerate to be k-duplications of the query point itself.

Parameters

• points – tensor of shape [N, P1, dims].

• query_points – tensor of shape [N, P2, dims]

• k – Integer.

• valid_num – tensor of shape [N,]

• max_distance – float representing the maximum distance that each neighbor can be. If there are no points within the distance, then the closest point is returned (regardless of distance). If this is set to None, then max_distance is not used.

Returns

• indices: tensor of shape [N, P2, k].

• padding: tensor of shape [N, P2 ,k] where 1 represents a padded point, and 0 represents an unpadded (real) point.

Return type

A pair of tensors

lingvo.tasks.car.car_lib.NeighborhoodIndices(points, query_points, k, points_padding=None, max_distance=None, sample_neighbors_uniformly=False)

Get indices to k-neighbors of query_points in points.

Padding is returned along-side indices. Non-padded points are guaranteed to be unique (non-repeated) points from original non-padded points.

Padded points arise due to either a lack of points (k exceeds the number of original non-padded points) or points are too far away (exceeds max distance).

Note: Padded point indices may refer to padded points from the original, or may be duplicates of the closest point.

TODO(weihan,jngiam): PointCNN implementation makes an assumption that padded points are repeated points from the original points. This behavior is maintained here, but we should update PointCNN to respect indices paddings.

Parameters

• points – tensor of shape [N, P1, dims].

• query_points – tensor of shape [N, P2, dims]

• k – Integer.

• points_padding – optional tensor of shape [N, P1] containing True/1.0 iff the point is a padded point. if None, then all points are considered real points.

• max_distance – float representing the maximum distance that each neighbor can be. If there are no points within the distance, then the closest point is returned (regardless of distance). If this is set to None, then no filtering by distance is performed.

• sample_neighbors_uniformly – boolean specifying whether to sample neighbors uniformly if they are within max distance.

Returns

• indices: tensor of shape [N, P2, k].

• padding: tensor of shape [N, P2, k] where 1 represents a padded point, and 0 represents an unpadded (real) point.

Return type

A pair of tensors

lingvo.tasks.car.car_lib.MatmulGather(source, indices)

Drop in replacement for tf.gather_nd() optimized for speed on TPU.

TODO(weihan): tf.gather_nd() is supposed to be implemented in the same way on TPU. Investigate why it’s much slower.

Parameters

• source – tensor of shape [N, P1, C]

• indices – tensor of shape [N, P2, K]

Returns

tensor of shape [N, P2, K, C]

lingvo.tasks.car.car_lib.FarthestPointSampler(points, padding, num_sampled_points, precomputed_squared_distance=None, num_seeded_points=0, random_seed=None)

Samples num_sampled_points from points using farthest point sampling.

Algorithm: 1. Start by selecting a random point and adding to a selected set. 2. For all remaining points, find the furthest point from those selected. 3. Add furthest point to selected. 4. Repeat 2-3 until num_sampled_points are selected.

More details at https://en.wikipedia.org/wiki/Farthest-first_traversal

This output of this function can be used with tf.array_ops.batch_gather to extract the desired points, for example: tf.array_ops.batch_gather(points, sampled_idx)

Parameters

• points – floating point tf.Tensor of shape [N, P1, dims]

• padding – A floating point tf.Tensor of shape [N, P1] with 0 if the point is real, and 1 otherwise.

• num_sampled_points – integer number of points to sample.

• precomputed_squared_distance – optional tf.Tensor of shape [N, P1, P1] of distances between each point. if None, distances will be computed on the fly.

• num_seeded_points – If num_seeded_points > 0, then the first num_seeded_points in points are considered to be seeded in the FPS sampling. Note that we assume that these points are not padded, and do not check padding when seeding them.

• random_seed – optional integer random seed to use with all the random ops.

Returns

A tuple of tf.Tensors (sampled_idx, closest_idx) of types (tf.int32, tf.int32).

sampled_idx is of shape [N, num_sampled_points] representing the indices selected using the sampler. This will have range of [0, P1].

closest_idx is of shape [N, P1] representing the indices of the closest sampled points for each input point. closest_idx is used in PCNN as part of the pooling operation: each point is assigned to the closest sampled point and a max is taken over them. This will have a range of [0, P2] with the index of the closest sampled point that remains.

lingvo.tasks.car.car_lib.MaxPool3D(points, point_features, pooling_idx, closest_idx)

Apply max pooling to a point cloud with computed sampling indices.

sampled_idx and closest_idx are the outputs of a sampler such as FurthestPointSampler.

The pooling operation results in a point cloud with fewer points, where the pooled points are specified by pooling_idx. Each element of pooling_idx contains an integer in the range [0, P1) containing the index of the point in points/points_features.

Max pooling is performed by assigning each point to its closest pooled point, and then taking a max over the features of points assigned. We assume that this mapping is provided by closest_idx, where each element should contain an integer in the range [0, P2) containing the index of the pooled point that each point is assigned to.

Note: This logic for pooling assumes that there will be at least one value > 0 per sampled region for each feature, otherwise it will return 0. Additionally, it does a reduce over a masked version of the features, so mean and min would not work without a change in the logic.

Parameters

• points – a floating point tf.Tensor with shape [N, P1, 3]

• point_features – a floating point tf.Tensor with shape [N, P1, C]

• pooling_idx – A tf.int32 tf.Tensor of shape [N, P2] with the index of which points we want to keep. Each value should be in the range [0, P1].

• closest_idx – A tf.int32 tf.Tensor of shape [N, P1] representing which sampled point is closest to each original point. Each value should be in the range of [0, P2].

Returns

A tuple of tf.Tensors (pooled_points, pooled_features).

pooled_points has shape [N, P2, 3] representing the locations of each selected point. P2 corresponds to num_pooled_points.

pooled_features has shape [N, P2, C] representing the pooled features at each point.

lingvo.tasks.car.car_lib.SegmentPool3D(points, point_features, pooling_idx, closest_idx, pooling_method='max')

Performs {min/max/average} pooling over a pointcloud given indices.

This should be functionally identical when using max to the above MaxPool3D function, except it turns out to be much more memory efficient on a TPU, and supports min/max/mean.

Parameters

• points – A float tf.Tensor of shape [N, P1, 3] with point locations.

• point_features – A float tf.Tensor of shape [N, P1, C] with point features.

• pooling_idx – A tf.int32 tf.Tensor of shape [N, P2] with the index of which points we want to keep. Each value should be in the range [0, P1].

• closest_idx – A tf.int32 tf.Tensor of shape [N, P1] representing which sampled point is closest to each original point. Each value should be in the range of [0, P2].

• pooling_method – A string for which pooling function to use. Should be one of {‘min’, ‘max’, ‘mean’}.

Returns

A float tf.Tensor of shape [N, P2, 3] with the pooled

point locations.

pooled_features: A float tf.Tensor of shape [N, P2, C] with the pooled

features.

Return type

pooled_points

Raises

ValueError – If pooling_method is not one of {min/max/mean}.