Point Transformer

Point Transformer classification and segmentation models.

First page of Point Transformer

2012.09164 · December 2020

Classes:

PointTransformerConv –

The Point Transformer layer from the
PointTransformerBlock –

Residual bottleneck around a PointTransformerConv: a linear projection down, vector attention, and a
PointTransformerTransitionDown –

Downsamples the cloud with farthest point sampling, then pools each centroid's \(k\)-NN neighborhood.
PointTransformerTransitionUp –

Upsamples features to the skip resolution by 3-NN interpolation and adds the projected skip features.
PointTransformerEncoderBlock –

One encoder stage: an optional transition down, then depth PointTransformerBlock units sharing a
PointTransformerDecoderBlock –

One decoder stage: an optional transition up onto the skip resolution, then depth
PointTransformerEncoder –

Stack of PointTransformerEncoderBlock stages, each but the first preceded by a transition down.
PointTransformerDecoder –

Stack of PointTransformerDecoderBlock stages that consume the encoder intermediates in reverse,
PointTransformerClassification –

Point Transformer classification model from the paper
PointTransformerSegmentation –

Point Transformer segmentation model from the paper

PointTransformerConv ¶

PointTransformerConv(
    in_channels: Union[int, Tuple[int, int]],
    out_channels: int,
    spatial_dim: int = 3,
    num_groups: int = 8,
    pos_nn: Optional[Callable[[Tensor], Tensor]] = None,
    attn_nn: Optional[Callable[[Tensor], Tensor]] = None,
    add_self_loops: bool = False,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

The Point Transformer layer from the "Point Transformer" paper by Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, Vladlen Koltun.

Note

This implementation was adapted from the PyTorch Geometric library, and supports the num_groups parameter to behave like the original implementation.

\[ \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i) \cup \{ i \}} \alpha_{i,j} \left(\mathbf{W}_3 \mathbf{x}_j + \delta_{ij} \right), \]

where the attention coefficients \(\alpha_{i,j}\) and positional embedding \(\delta_{ij}\) are computed as

\[ \alpha_{i,j}= \textrm{softmax} \left( \gamma_\mathbf{\Theta} (\mathbf{W}_1 \mathbf{x}_i - \mathbf{W}_2 \mathbf{x}_j + \delta_{i,j}) \right) \]

and

\[ \delta_{i,j}= h_{\mathbf{\Theta}}(\mathbf{p}_i - \mathbf{p}_j), \]

with \(\gamma_\mathbf{\Theta}\) and \(h_\mathbf{\Theta}\) denoting neural networks, i.e. MLPs, and \(\mathbf{P} \in \mathbb{R}^{N \times D}\) defines the position of each point.

Parameters:

in_channels (int or tuple) –

Size of each input sample, or -1 to derive the size from the first input(s) to the forward method. A tuple corresponds to the sizes of source and target dimensionalities.
out_channels (int) –

Size of each output sample.
pos_nn (Module, default: None ) –

A neural network \(h_\mathbf{\Theta}\) which maps relative spatial coordinates pos_j - pos_i of shape \([-1, 3]\) to shape \([-1, \text{out\_channels}]\). Will default to a torch.nn.Linear transformation if not further specified.
attn_nn (Module, default: None ) –

A neural network \(\gamma_\mathbf{\Theta}\) which maps transformed node features of shape \([-1, \text{out\_channels}]\) to shape \([-1, \text{out\_channels}]\).
add_self_loops (bool, default: False ) –

If False, do not add self-loops to the input graph.

Shapes

input: node features \((|\mathcal{V}|, F_{in})\) or \(((|\mathcal{V_s}|, F_{s}), (|\mathcal{V_t}|, F_{t}))\) if bipartite, positions \((|\mathcal{V}|, 3)\) or \(((|\mathcal{V_s}|, 3), (|\mathcal{V_t}|, 3))\) if bipartite, edge indices \((2, |\mathcal{E}|)\)
output: node features \((|\mathcal{V}|, F_{out})\) or \(((|\mathcal{V}_t|, F_{out}))\) if bipartite

PointTransformerBlock ¶

PointTransformerBlock(
    in_channels: int,
    out_channels: int,
    spatial_dim: int = 3,
    num_groups: int = 8,
    add_self_loops: bool = False,
    act: Union[str, Callable, None] = "relu",
    act_first: bool = False,
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Residual bottleneck around a PointTransformerConv: a linear projection down, vector attention, and a linear projection back, each followed by normalization and activation.

PointTransformerTransitionDown ¶

PointTransformerTransitionDown(
    in_channels: int,
    out_channels: int,
    num_neighbors: int = 16,
    ratio: float = 0.25,
    spatial_dim: int = 3,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    pool: str = "max",
)

Bases: Module

Downsamples the cloud with farthest point sampling, then pools each centroid's \(k\)-NN neighborhood.

The MLP sees the relative position of a neighbor concatenated with its features, and the neighborhood is reduced with pool.

PointTransformerTransitionUp ¶

PointTransformerTransitionUp(
    in_channels: int,
    out_channels: int,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Upsamples features to the skip resolution by 3-NN interpolation and adds the projected skip features.

PointTransformerEncoderBlock ¶

PointTransformerEncoderBlock(
    channels: int,
    depth: int,
    num_groups: int,
    num_neighbors: int,
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    downsample: Optional[Module] = None,
)

Bases: Module

One encoder stage: an optional transition down, then depth PointTransformerBlock units sharing a single \(k\)-NN graph built on the stage's own resolution.

PointTransformerDecoderBlock ¶

PointTransformerDecoderBlock(
    channels: int,
    depth: int,
    num_groups: int,
    num_neighbors: int,
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    upsample: Optional[Module] = None,
)

Bases: Module

One decoder stage: an optional transition up onto the skip resolution, then depth PointTransformerBlock units sharing a single \(k\)-NN graph built on that resolution.

PointTransformerEncoder ¶

PointTransformerEncoder(
    channels: Sequence[int],
    depths: Sequence[int],
    num_groups: Sequence[int],
    num_neighbors: Sequence[int],
    ratios: Sequence[float],
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Stack of PointTransformerEncoderBlock stages, each but the first preceded by a transition down.

When return_intermediates=True is passed to forward, the input features and point cloud of every stage are returned in fine-to-coarse order, ready to be consumed as decoder skips.

PointTransformerDecoder ¶

PointTransformerDecoder(
    channels: Sequence[int],
    depths: Sequence[int],
    num_groups: Sequence[int],
    num_neighbors: Sequence[int],
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    upsample: Optional[Module] = None,
)

Bases: Module

Stack of PointTransformerDecoderBlock stages that consume the encoder intermediates in reverse, walking the features back to full resolution.

PointTransformerClassification ¶

PointTransformerClassification(
    in_channels: int,
    num_classes: int,
    *,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_num_groups: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    ratios: Sequence[float],
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    global_pool: PoolLike = "max",
    dropout: float = 0.0,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: ClassificationModel

Point Transformer classification model from the paper Point Transformer by Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, Vladlen Koltun.

A hierarchical encoder of vector-attention PointTransformerBlock stages interleaved with PointTransformerTransitionDown downsampling, followed by global pooling and a linear head.

Parameters:

in_channels (int) –

Number of input feature channels. Pass \(0\) to use the raw positions as features.
num_classes (int) –

Number of output classes.
encoder_channels (Sequence[int]) –

Feature width of each encoder stage.
encoder_depths (Sequence[int]) –

Number of PointTransformerBlock blocks per encoder stage.
encoder_num_groups (Sequence[int]) –

Number of shared-weight vector-attention groups per encoder stage.
encoder_num_neighbors (Sequence[int]) –

Number of neighbors in the \(k\)-NN graph of each encoder stage.
ratios (Sequence[float]) –

Farthest-point-sampling keep ratio for each downsampling transition (length one less than the number of encoder stages).
spatial_dim (int, default: 3 ) –

Dimensionality of the point coordinates.
add_self_loops (bool, default: False ) –

Whether to add self-loops to each neighborhood graph.
global_pool (PoolLike, default: 'max' ) –

Pooling used to aggregate point features into a per-cloud vector.
dropout (float, default: 0.0 ) –

Dropout probability applied before the classification head.
act (Union[str, Callable, None], default: 'relu' ) –

Activation used across the network.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the activation factory.
act_first (bool, default: False ) –

Whether to apply the activation before the normalization.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization used across the network.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the normalization factory.

Shape

x: \((N, \text{in\_channels})\), or None to fall back to pos.
pos: \((N, 3)\) point coordinates.
batch: \((N,)\) per-point batch index.
output: \((B, \text{num\_classes})\) class logits.

Methods:

configure_embeddings –

Build the embedding MLP lifting the input features to the first encoder width.
configure_encoder –

Build the PointTransformerEncoder backbone.

Attributes:

num_features (int) –

Feature dimension \(C\) of the encoder output.

num_features `property` ¶

num_features: int

Feature dimension \(C\) of the encoder output.

configure_embeddings ¶

configure_embeddings() -> MLP

Build the embedding MLP lifting the input features to the first encoder width.

configure_encoder ¶

configure_encoder() -> PointTransformerEncoder

Build the PointTransformerEncoder backbone.

PointTransformerSegmentation ¶

PointTransformerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_num_groups: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_num_groups: Sequence[int],
    decoder_num_neighbors: Sequence[int],
    ratios: Sequence[float],
    spatial_dim: int = 3,
    add_self_loops: bool = False,
    dropout: float = 0.0,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SegmentationModel

Point Transformer segmentation model from the paper Point Transformer by Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, Vladlen Koltun.

An encoder-decoder with skip connections: vector-attention PointTransformerBlock stages with PointTransformerTransitionDown downsampling, mirrored by PointTransformerTransitionUp upsampling, followed by a per-point linear head.

Parameters:

in_channels (int) –

Number of input feature channels. Pass \(0\) to use the raw positions as features.
num_classes (int) –

Number of output classes.
encoder_channels (Sequence[int]) –

Feature width of each encoder stage.
encoder_depths (Sequence[int]) –

Number of PointTransformerBlock blocks per encoder stage.
encoder_num_groups (Sequence[int]) –

Number of shared-weight vector-attention groups per encoder stage.
encoder_num_neighbors (Sequence[int]) –

Number of neighbors in the \(k\)-NN graph of each encoder stage.
decoder_channels (Sequence[int]) –

Feature width of each decoder stage (the last entry is the head width).
decoder_depths (Sequence[int]) –

Number of PointTransformerBlock blocks per decoder stage.
decoder_num_groups (Sequence[int]) –

Number of shared-weight vector-attention groups per decoder stage.
decoder_num_neighbors (Sequence[int]) –

Number of neighbors in the \(k\)-NN graph of each decoder stage.
ratios (Sequence[float]) –

Farthest-point-sampling keep ratio for each downsampling transition (length one less than the number of encoder stages).
spatial_dim (int, default: 3 ) –

Dimensionality of the point coordinates.
add_self_loops (bool, default: False ) –

Whether to add self-loops to each neighborhood graph.
dropout (float, default: 0.0 ) –

Dropout probability applied to the per-point features before the head.
act (Union[str, Callable, None], default: 'relu' ) –

Activation used across the network.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the activation factory.
act_first (bool, default: False ) –

Whether to apply the activation before the normalization.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization used across the network.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Optional keyword arguments for the normalization factory.

Shape

x: \((N, \text{in\_channels})\), or None to fall back to pos.
pos: \((N, 3)\) point coordinates.
batch: \((N,)\) per-point batch index.
output: \((N, \text{num\_classes})\) per-point class logits.

Methods:

configure_embeddings –

Build the embedding MLP lifting the input features to the first encoder width.
configure_encoder –

Build the PointTransformerEncoder backbone.
configure_decoder –

Build the PointTransformerDecoder upsampling the coarsest features back through the encoder skips.

Attributes:

num_features (int) –

Feature dimension \(C\) of the decoder output.

num_features `property` ¶

num_features: int

Feature dimension \(C\) of the decoder output.

configure_embeddings ¶

configure_embeddings() -> MLP

Build the embedding MLP lifting the input features to the first encoder width.

configure_encoder ¶

configure_encoder() -> PointTransformerEncoder

Build the PointTransformerEncoder backbone.

configure_decoder ¶

configure_decoder() -> PointTransformerDecoder

Build the PointTransformerDecoder upsampling the coarsest features back through the encoder skips.

Point Transformer

PointTransformerConv ¶

PointTransformerBlock ¶

PointTransformerTransitionDown ¶

PointTransformerTransitionUp ¶

PointTransformerEncoderBlock ¶

PointTransformerDecoderBlock ¶

PointTransformerEncoder ¶

PointTransformerDecoder ¶

PointTransformerClassification ¶

num_features property ¶

configure_embeddings ¶

configure_encoder ¶

PointTransformerSegmentation ¶

num_features property ¶

configure_embeddings ¶

configure_encoder ¶

configure_decoder ¶

num_features `property` ¶

num_features `property` ¶