grid_pool

Grid-based point cloud pooling (downsampling).

Clusters points by quantized grid coordinates and reduces features via scatter operations. This is an alternative to code-space pooling (SerializedPooling) used in PTV3 Mode 2 (Sonata) and Mode 3 (Utonia).

Classes:

GridPool –

Grid-based downsampling that clusters points by quantized coordinates.

GridPool ¶

GridPool(
    in_channels: int,
    out_channels: int,
    stride: int = 2,
    bias: bool = True,
    act: Union[str, Callable, None] = None,
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    reduce: str = "max",
)

Bases: Module

Grid-based downsampling that clusters points by quantized coordinates.

Divides pos_grid by stride, groups unique voxels via torch.unique, and reduces features with torch_scatter.segment_csr.

Parameters:

in_channels (int) –

Number of input feature channels.
out_channels (int) –

Number of output feature channels.
stride (int, default: 2 ) –

Spatial stride for grid quantization.
act (Union[str, Callable, None], default: None ) –

Activation layer applied after projection.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the activation function.
act_first (bool, default: False ) –

Apply activation before normalization.
norm (Union[str, Callable, None], default: None ) –

Normalization layer applied after projection.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the normalization layer.
reduce (str, default: 'max' ) –

Scatter reduction ("max", "mean", "sum", "min").

Methods:

forward –

Downsample points by grid quantization.

forward ¶

forward(
    x: Tensor,
    pos_grid: Tensor,
    batch: Tensor,
    pos: Optional[Tensor] = None,
    condition: Optional[str] = None,
) -> Tuple[
    Tensor, Tensor, Tensor, Tensor, Optional[Tensor]
]

Downsample points by grid quantization.

Parameters:

x (Tensor) –

Point features of shape \((N, C_{in})\).
pos_grid (Tensor) –

Integer grid coordinates of shape \((N, 3)\). Must be non-negative and pool to coordinates below \(2^{16}\): batch index and coordinates are bit-packed into a single int64 cluster key, so out-of-range coordinates would merge clusters across samples.
batch (Tensor) –

Batch indices of shape \((N,)\).
pos (Optional[Tensor], default: None ) –

Optional real-valued positions of shape \((N, 3)\) to mean-pool alongside the features (e.g. for downstream rotary position embedding).
condition (Optional[str], default: None ) –

Optional condition name selecting the inner norm when the norm is a PDNorm.

Returns:

Tensor –

Tuple of (x_pooled, pos_grid_pooled, batch_pooled, pooling_inverse, pos_pooled).
Tensor –

pooling_inverse maps each input point to its pooled cluster index.
Tensor –

pos_pooled is None when pos is not provided.