Skip to content

OctFormer

OctFormer classification and segmentation models.

First page of OctFormer: Octree-based Transformers for 3D Point Clouds

2305.03045 · May 2023

Classes:

  • CPE –

    Conditional positional encoding: a depthwise octree convolution followed by batch normalization.

  • OctFormerBlock –

    Transformer block over octree patches: a CPE residual, octree attention and an MLP, both pre-normed.

  • OctFormerEncoderLayer –

    One encoder stage: an optional octree convolution downsampling followed by num_blocks OctFormerBlock units.

  • OctFormerEncoder –

    Stack of OctFormerEncoderLayer stages, each running one octree depth coarser than the previous one.

  • OctFormerDecoder –

    Feature pyramid decoder: projects every encoder stage to fpn_channels, merges them top-down, and

  • OctreePatchEmbed –

    Convolutional stem that embeds the octree signal, halving the resolution once per channel step.

  • OctFormerClassification –

    OctFormer classification model from

  • OctFormerSegmentation –

    OctFormer segmentation model from

CPE 

CPE(
    in_channels: int,
    kernel_size: Union[int, Sequence[int]] = 3,
    nempty: bool = False,
    bias: bool = False,
    use_dwconv: bool = False,
)

Bases: Module

Conditional positional encoding: a depthwise octree convolution followed by batch normalization.

OctFormerBlock 

OctFormerBlock(
    channels: int,
    num_heads: int,
    patch_size: int = 32,
    dilation: int = 0,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.0,
    nempty: bool = False,
    use_rpe: bool = True,
    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,
    bias: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
)

Bases: Module

Transformer block over octree patches: a CPE residual, octree attention and an MLP, both pre-normed.

Note

cpe_first places the CPE residual before the attention; False places it after the attention residual.

OctFormerEncoderLayer 

OctFormerEncoderLayer(
    channels: int,
    num_heads: int,
    patch_size: int = 32,
    dilation: int = 0,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: Union[float, Sequence[float]] = 0.0,
    nempty: bool = False,
    use_checkpoint: bool = True,
    use_rpe: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
    num_blocks: int = 2,
    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,
    bias: bool = True,
    downsample: Optional[Module] = None,
)

Bases: Module

One encoder stage: an optional octree convolution downsampling followed by num_blocks OctFormerBlock units.

Blocks alternate between a dilation of 1 and the configured dilation, so consecutive blocks attend to neighboring and to spread-out patches in turn.

OctFormerEncoder 

OctFormerEncoder(
    channels: Sequence[int],
    num_blocks: Sequence[int],
    num_heads: Sequence[int],
    patch_size: int = 26,
    dilation: int = 4,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.5,
    nempty: bool = False,
    use_checkpoint: bool = True,
    use_rpe: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
    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,
    bias: bool = True,
)

Bases: Module

Stack of OctFormerEncoderLayer stages, each running one octree depth coarser than the previous one.

When return_intermediates=True is passed to forward, the input features of every stage but the first are returned in coarse-to-fine order, ready to be consumed by OctFormerDecoder.

OctFormerDecoder 

OctFormerDecoder(
    channels: Sequence[int],
    fpn_channels: int,
    num_ups: int = 1,
    nempty: bool = True,
    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,
    bias: bool = True,
)

Bases: Module

Feature pyramid decoder: projects every encoder stage to fpn_channels, merges them top-down, and sums the results upsampled to the finest depth. num_ups octree deconvolutions then undo the stem strides.

OctreePatchEmbed 

OctreePatchEmbed(
    channels: Sequence[int],
    nempty: 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,
    bias: bool = True,
)

Bases: Module

Convolutional stem that embeds the octree signal, halving the resolution once per channel step.

OctFormerClassification 

OctFormerClassification(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Union[int, Sequence[int]],
    encoder_channels: Sequence[int],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    num_blocks: Sequence[int],
    num_heads: Sequence[int],
    patch_size: int = 26,
    dilation: int = 4,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.5,
    nempty: bool = True,
    use_checkpoint: bool = True,
    use_rpe: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    head_act: Union[str, Callable, None] = "gelu",
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    dropout: float = 0.0,
    global_pool: PoolLike = "mean",
)

Bases: ClassificationModel

OctFormer classification model from OctFormer: Octree-based Transformers for 3D Point Clouds by Peng-Shuai Wang.

An octree convolution stem embeds the input signal, then attention runs over patches of equal point count sorted along the octree's space-filling curve, one stage per octree depth. Point features are pooled globally after the encoder for classification.

Methods:

  • configure_stem –

    Build the OctreePatchEmbed stem.

  • configure_encoder –

    Build the OctFormerEncoder backbone.

  • get_encoder_depth –

    Octree depth at which the first encoder stage runs, once the stem downsamplings are accounted for.

  • get_head_depth –

    Octree depth of the encoder output, which is the depth the head pools over.

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_stem 

configure_stem() -> Module

Build the OctreePatchEmbed stem.

configure_encoder 

configure_encoder() -> Module

Build the OctFormerEncoder backbone.

get_encoder_depth 

get_encoder_depth(depth: int) -> int

Octree depth at which the first encoder stage runs, once the stem downsamplings are accounted for.

Parameters:

  • depth (int) –

    Octree depth of the input signal.

Returns:

  • int –

    The octree depth of the encoder's first stage.

get_head_depth 

get_head_depth(depth: int) -> int

Octree depth of the encoder output, which is the depth the head pools over.

Parameters:

  • depth (int) –

    Octree depth of the input signal.

Returns:

  • int –

    The octree depth of the encoder's last stage.

OctFormerSegmentation 

OctFormerSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    stem_channels: Union[int, Sequence[int]],
    channels: Sequence[int],
    num_blocks: Sequence[int],
    num_heads: Sequence[int],
    head_channels: Optional[
        Union[int, Sequence[int]]
    ] = None,
    fpn_channels: int,
    patch_size: int = 26,
    dilation: int = 4,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = True,
    qk_scale: Optional[float] = None,
    attn_drop: float = 0.0,
    proj_drop: float = 0.0,
    drop_path: float = 0.5,
    nempty: bool = True,
    use_checkpoint: bool = True,
    use_rpe: bool = True,
    use_dwconv: bool = False,
    cpe_first: bool = True,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    act_first: bool = False,
    head_act: Union[str, Callable, None] = "gelu",
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = True,
    dropout: float = 0.5,
)

Bases: SegmentationModel

OctFormer segmentation model from OctFormer: Octree-based Transformers for 3D Point Clouds by Peng-Shuai Wang.

The octree attention encoder is followed by a feature pyramid decoder that merges every stage at fpn_channels and upsamples back to the input octree depth, then a per-point MLP head.

Methods:

Attributes:

  • num_features (int) –

    Feature dimension \(C\) of the features entering the head.

num_features property 

num_features: int

Feature dimension \(C\) of the features entering the head.

configure_stem 

configure_stem() -> Module

Build the OctreePatchEmbed stem.

configure_encoder 

configure_encoder() -> Module

Build the OctFormerEncoder backbone.

configure_decoder 

configure_decoder() -> Module

Build the OctFormerDecoder, with one upsampling per stem downsampling.