ops

Tensor operations on packed batches: safe division, softmax, voxel hashing, interpolation, and decimation.

Functions:

safe_divide –

Safely divide two tensors, returning a default value if the denominator is zero.
softmax –

Apply softmax on a packed x tensor.
voxel_grid_fnv –

FNV-1a 64-bit hash of integer voxel-grid coordinates. \((N, D) \to (N,)\).
first_permutation –

Index of the first occurrence of each cluster id in a consecutive cluster tensor.
knn_interpolate –

k-NN interpolation with inverse-distance weighting.
decimate_indices –

Decimate indices from a packed batch index tensor.
decimate –

Decimates each input tensor by the given factor.
pad_tail –

Pad the tail of a tensor with a fill value.
offset_index –

Offset per-element row indices into the packed row layout of a collated batch.

safe_divide ¶

safe_divide(
    a: Tensor,
    b: Tensor,
    /,
    default: Union[float, Tensor] = float("nan"),
) -> Tensor

Safely divide two tensors, returning a default value if the denominator is zero.

Note

If the inputs are not floating point numbers, they will be converted to floating point numbers (float32).

Parameters:

a (Tensor) –

The numerator tensor.
b (Tensor) –

The denominator tensor.
default (Union[float, Tensor], default: float('nan') ) –

The default value to return if the denominator is zero.

Returns:

Tensor –

The result of the division.

Example

>>> safe_divide(torch.tensor([1.0, 2.0, 3.0]), torch.tensor([1.0, 0.0, 1.0]))
tensor([1., nan, 3.])
>>> safe_divide(torch.tensor([1.0, 2.0, 3.0]), torch.tensor([1.0, 0.0, 1.0]), default=0.0)
tensor([1., 0., 3.])
>>> safe_divide(torch.tensor([1, 2, 3]), torch.tensor([1, 0, 1]), default=torch.tensor([0, 0, 0]))
tensor([1., 0., 3.])

softmax ¶

softmax(x: Tensor, batch: Tensor, dim: int = 0) -> Tensor

Apply softmax on a packed x tensor. The x tensor is expected to be of shape \((N, *)\), where \(N\) is the number of nodes and \(*\) is the feature size. The batch tensor must be of shape \((N,)\) and must be contiguous.

Note

This function is adapted from the torch_geometric package, and requires the torch-scatter package.

Parameters:

x (Tensor) –

The x tensor of shape \((N, *)\).
batch (Tensor) –

The batch tensor of shape \((N,)\).
dim (int, default: 0 ) –

The dimension along which to apply the softmax.

Returns:

Tensor –

The softmaxed tensor of shape \((N, *)\).

voxel_grid_fnv ¶

voxel_grid_fnv(
    pos: Tensor,
    size: float,
    start: Optional[Tensor] = None,
    *,
    return_inverse: Literal[False] = False,
    return_counts: Literal[False] = False,
) -> Tensor

voxel_grid_fnv(
    pos: Tensor,
    size: float,
    start: Optional[Tensor] = None,
    *,
    return_inverse: Literal[True],
    return_counts: Literal[False] = False,
) -> Tuple[Tensor, Tensor]

voxel_grid_fnv(
    pos: Tensor,
    size: float,
    start: Optional[Tensor] = None,
    *,
    return_inverse: Literal[False] = False,
    return_counts: Literal[True],
) -> Tuple[Tensor, Tensor]

voxel_grid_fnv(
    pos: Tensor,
    size: float,
    start: Optional[Tensor] = None,
    *,
    return_inverse: Literal[True],
    return_counts: Literal[True],
) -> Tuple[Tensor, Tensor, Tensor]

voxel_grid_fnv(
    pos: Tensor,
    size: float,
    start: Optional[Tensor] = None,
    *,
    return_inverse: bool = False,
    return_counts: bool = False,
) -> Union[Tensor, Tuple[Tensor, ...]]

FNV-1a 64-bit hash of integer voxel-grid coordinates. \((N, D) \to (N,)\).

Parameters:

pos (Tensor) –

Point positions of shape \((N, D)\).
size (float) –

Voxel side length in the same units as pos.
start (Optional[Tensor], default: None ) –

Optional voxel-grid origin. When None, the grid origin is implicit via the internal pos_grid -= pos_grid.min(0) shift.
return_inverse (bool, default: False ) –

If True, also return the per-point consecutive voxel index in \([0, V)\), following the semantics of torch.unique(..., return_inverse=True).
return_counts (bool, default: False ) –

If True, also return the per-voxel point count of shape \((V,)\).

Returns:

Union[Tensor, Tuple[Tensor, ...]] –

hashed of shape \((N,)\) when both flags are False. With return_inverse=True adds
Union[Tensor, Tuple[Tensor, ...]] –

inverse of shape \((N,)\); with return_counts=True adds count of shape \((V,)\); both
Union[Tensor, Tuple[Tensor, ...]] –

flags enabled returns (hashed, inverse, count).

first_permutation ¶

first_permutation(
    cluster: Tensor, num_clusters: Optional[int] = None
) -> Tensor

Index of the first occurrence of each cluster id in a consecutive cluster tensor.

The permutation returned by consecutive_cluster picks a backend-dependent representative per cluster (the last occurrence on CPU, a nondeterministic one on CUDA). This helper always picks the first occurrence, so tensor[first_permutation(cluster)] is deterministic across devices.

Parameters:

cluster (Tensor) –

Consecutive cluster indices of shape \((N,)\) with values in \([0, V)\).
num_clusters (Optional[int], default: None ) –

Number of clusters \(V\). Inferred as cluster.max() + 1 when None.

Returns:

Tensor –

Long tensor of shape \((V,)\) holding, per cluster id, the smallest index in cluster with that id.

Example

>>> cluster = torch.tensor([1, 0, 1, 2, 0])
>>> first_permutation(cluster)
tensor([1, 0, 3])

knn_interpolate ¶

knn_interpolate(
    x: Tensor,
    pos_x: Tensor,
    pos_y: Tensor,
    batch_x: OptTensor = None,
    batch_y: OptTensor = None,
    k: int = 3,
    num_workers: int = 1,
    weighting: Literal["squared", "inverse"] = "squared",
    eps: float = 1e-16,
) -> Tensor

k-NN interpolation with inverse-distance weighting.

From PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space.

For each point \(y\) with position \(\mathbf{p}(y)\), its interpolated features \(\mathbf{f}(y)\) are given by

\[ \mathbf{f}(y) = \frac{\sum_{i=1}^k w(x_i) \mathbf{f}(x_i)}{\sum_{i=1}^k w(x_i)} \]

where \(\{ x_1, \ldots, x_k \}\) are the \(k\) nearest points to \(y\) and the weights \(w(x_i)\) depend on the chosen weighting scheme:

"squared" (default, torch_geometric convention): \(w(x_i) = 1 / d(\mathbf{p}(y), \mathbf{p}(x_i))^2\)
"inverse" (PointNet++ three_interpolation convention): \(w(x_i) = 1 / d(\mathbf{p}(y), \mathbf{p}(x_i))\)

Note

Adapted from the torch_geometric package. Requires torch-cluster.

Parameters:

x (Tensor) –

Node feature matrix \(\mathbf{X} \in \mathbb{R}^{N \times F}\).
pos_x (Tensor) –

Node position matrix \(\in \mathbb{R}^{N \times d}\).
pos_y (Tensor) –

Upsampled node position matrix \(\in \mathbb{R}^{M \times d}\).
batch_x (OptTensor, default: None ) –

Batch vector \(\mathbf{b_x} \in \{ 0, \ldots, B-1 \}^N\), assigning each node from \(\mathbf{X}\) to a specific example.
batch_y (OptTensor, default: None ) –

Batch vector \(\mathbf{b_y} \in \{ 0, \ldots, B-1 \}^M\), assigning each node from \(\mathbf{Y}\) to a specific example.
k (int, default: 3 ) –

Number of neighbors.
num_workers (int, default: 1 ) –

Number of workers for computation. Has no effect when batch_x or batch_y is not None, or the input lies on GPU.
weighting (Literal['squared', 'inverse'], default: 'squared' ) –

Weighting scheme for neighbors. "squared" for \(1/d^2\) weights (torch_geometric default) or "inverse" for \(1/d\) weights (PointNet++ convention).
eps (float, default: 1e-16 ) –

Small value to avoid division by zero.

Returns:

Tensor –

Interpolated features \(\in \mathbb{R}^{M \times F}\).

decimate_indices ¶

decimate_indices(
    batch: Tensor,
    factor: float,
    generator: Optional[Generator] = None,
) -> Tuple[Tensor, Tensor]

Decimate indices from a packed batch index tensor. This function will return the decimated indices by a given factor along with the decimated batch indices.

Note

This function is similar to the decimation_indices function in the torch-geometric package, except that this function uses the batch tensor instead of the ptr tensor representation.

Parameters:

batch (Tensor) –

The packed batch index tensor.
factor (float) –

The factor to decimate the indices by.
generator (Optional[Generator], default: None ) –

The generator to use for the random permutation.

Returns:

Tuple[Tensor, Tensor] –

The decimated indices and the decimated batch indices.

Examples:

>>> batch = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3])
>>> decimate_indices(batch, 2)  # doctest: +SKIP
(tensor([ 0,  4,  7,  6,  9, 10]), tensor([0, 1, 2, 2, 3, 3]))

decimate ¶

decimate(
    tensors: Tuple[Tensor, ...],
    batch: Tensor,
    factor: int,
    generator: Optional[Generator] = None,
) -> Tuple[Tuple[Tensor, ...], Tensor]

Decimates each input tensor by the given factor. This will return the decimated tensors along with the decimated batch indices.

Note

This function is similar to the decimate function introduced in the torch-geometric RandLANet example.

Parameters:

tensors (Tuple[Tensor, ...]) –

A tuple of tensors to decimate.
batch (Tensor) –

The batch tensor of shape \((N,)\).
factor (int) –

The factor to decimate the tensors by.
generator (Optional[Generator], default: None ) –

The generator to use for the random permutation.

Returns:

Tuple[Tuple[Tensor, ...], Tensor] –

A tuple of decimated tensors and the decimated batch indices.

Examples:

>>> tensors = (torch.randn(10, 3), torch.randn(10, 4))
>>> batch = torch.tensor([0, 1, 1, 1, 2, 2, 2, 2, 3, 3])
>>> decimate(tensors, batch, 2)  # doctest: +SKIP
((tensor([[-1.4570, -0.1023, -0.5992],
        [ 0.2408,  0.1325,  0.7642],
        [-0.2104, -1.4391,  0.5214],
        [ 1.6192,  1.4506,  0.2695],
        [ 0.3488,  0.9676, -0.4657]]),
tensor([[-0.1933,  0.6526, -1.9006,  0.2286],
        [ 1.2888,  0.0523, -1.5469,  0.7567],
        [ 0.9442, -0.1849,  1.0608,  0.2083],
        [ 0.4788,  1.3537, -0.1593, -0.4249],
        [ 1.3065,  0.4598,  0.2618, -0.7599]])),
tensor([0, 1, 2, 2, 3]))

pad_tail ¶

pad_tail(
    tensor: Tensor,
    pad_size: int,
    dim: int,
    fill_value: float = 0,
) -> Tensor

Pad the tail of a tensor with a fill value.

Parameters:

tensor (Tensor) –

The tensor to pad.
pad_size (int) –

The size of the padding that will be added to the tail of the tensor.
dim (int) –

The dimension along which to pad the tensor.
fill_value (float, default: 0 ) –

The value to fill the padding with.

Returns:

Tensor –

The padded tensor.

Examples:

>>> tensor = torch.tensor([1, 2, 3])
>>> pad_tail(tensor, pad_size=2, dim=0, fill_value=0)
tensor([1, 2, 3, 0, 0])

>>> tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> pad_tail(tensor, pad_size=2, dim=0, fill_value=0)
tensor([[1, 2, 3],
        [4, 5, 6],
        [0, 0, 0],
        [0, 0, 0]])

>>> tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
>>> pad_tail(tensor, pad_size=2, dim=1, fill_value=0)
tensor([[1, 2, 3, 0, 0],
        [4, 5, 6, 0, 0]])

offset_index ¶

offset_index(
    index: Tensor, index_batch: Tensor, batch: Tensor
) -> Tensor

Offset per-element row indices into the packed row layout of a collated batch.

collate concatenates a row map such as inverse or index as-is, so the entries of each batch element still address that element's own rows. This shifts every entry by the number of rows of the elements collated before it.

Parameters:

index (Tensor) –

Per-element row indices.
index_batch (Tensor) –

Batch index of each entry of index, the batch_<key> tensor collate emits for cat_keys.
batch (Tensor) –

Batch index of the rows index addresses (batch for an inverse map, batch_origin_pos for an index map).

Returns:

Tensor –

The offset row indices.

Shape

index: \((M,)\)
index_batch: \((M,)\)
batch: \((N,)\)
output: \((M,)\)

Example

import torch
from torch_pointcloud.utils.ops import offset_index

inverse = torch.tensor([0, 1, 1, 0, 2, 2, 1])
batch_inverse = torch.tensor([0, 0, 0, 1, 1, 1, 1])
batch = torch.tensor([0, 0, 1, 1, 1])
offset_index(inverse, batch_inverse, batch)  # tensor([0, 1, 1, 2, 4, 4, 3])