Skip to content

point_transformer

Point Transformer classification and segmentation models.

First page of Point Transformer

2012.09164 · December 2020

Classes:

PointTransformerBlock [source] 

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.

The attention MLP normalizes and activates the attention logits before its first linear layer, and the query, key and value projections carry a bias.

PointTransformerTransitionDown [source] 

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 [source] 

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.

PointTransformerGlobalTransition [source] 

PointTransformerGlobalTransition(
    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

Coarsest-level transition of the segmentation decoder: appends a projected global mean feature to every point and projects the pair back to channels.

PointTransformerEncoderBlock [source] 

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 [source] 

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 [source] 

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 [source] 

PointTransformerDecoder(
    channels: Sequence[int],
    depths: Sequence[int],
    num_groups: Sequence[int],
    num_neighbors: Sequence[int],
    global_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

A PointTransformerGlobalTransition and one PointTransformerBlock at the coarsest resolution, then a stack of PointTransformerDecoderBlock stages that consume the encoder intermediates in reverse, walking the features back to full resolution.

Parameters:

  • channels (Sequence[int]) –

    Feature width entering the decoder followed by the width of each upsampling stage.

  • depths (Sequence[int]) –

    Number of PointTransformerBlock blocks per upsampling stage.

  • num_groups (Sequence[int]) –

    Number of shared-weight vector-attention groups per upsampling stage.

  • num_neighbors (Sequence[int]) –

    Number of neighbors in the \(k\)-NN graph of each upsampling stage.

  • global_num_neighbors (int) –

    Number of neighbors in the \(k\)-NN graph of the coarsest-level block.

PointTransformerClassification [source] 

PointTransformerClassification(
    in_channels: int,
    num_classes: int,
    *,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_attention_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_attention_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:

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 [source] 

configure_embeddings() -> MLP

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

configure_encoder [source] 

configure_encoder() -> PointTransformerEncoder

Build the PointTransformerEncoder backbone.

PointTransformerSegmentation [source] 

PointTransformerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_attention_groups: Sequence[int],
    encoder_num_neighbors: Sequence[int],
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_attention_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: SemanticSegmentationModel

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, a PointTransformerGlobalTransition mixing a global mean feature into the coarsest stage, PointTransformerTransitionUp upsampling mirroring the encoder, and a per-point MLP 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_attention_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_attention_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 [source] 

configure_embeddings() -> MLP

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

configure_encoder [source] 

configure_encoder() -> PointTransformerEncoder

Build the PointTransformerEncoder backbone.

configure_decoder [source] 

configure_decoder() -> PointTransformerDecoder

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