Skip to content

SPFormer-UNet

SPFormer-UNet segmentation model.

First page of Superpoint Transformer for 3D Scene Instance Segmentation

2211.15766 · November 2022

Classes:

SPFormerUNetEncoderBlock 

SPFormerUNetEncoderBlock(
    channels: int,
    depth: int,
    indice_key: str,
    *,
    downsample: Optional[Module] = None,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

One encoder stage: an optional stride-2 downsample followed by depth residual blocks.

Parameters:

  • channels (int) –

    Channel width of this stage (the residual blocks run at this width).

  • depth (int) –

    Number of residual blocks.

  • indice_key (str) –

    SpConv submanifold index key shared by the residual blocks.

  • downsample (Optional[Module], default: None ) –

    Stride-2 down-conv applied before the blocks, or None for the first (full-resolution) stage.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation passed to create_act.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization passed to create_norm.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the normalization.

SPFormerUNetDecoderBlock 

SPFormerUNetDecoderBlock(
    channels: int,
    depth: int,
    indice_key: str,
    *,
    upsample: Module,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

One decoder stage: upsample, concatenate the encoder skip, then depth residual blocks.

Parameters:

  • channels (int) –

    Output channel width of this stage.

  • depth (int) –

    Number of residual blocks.

  • indice_key (str) –

    SpConv submanifold index key shared by the residual blocks.

  • upsample (Module) –

    Inverse conv mapping the deeper feature to channels.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation passed to create_act.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization passed to create_norm.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the normalization.

SPFormerUNetEncoder 

SPFormerUNetEncoder(
    in_channels: int,
    channels: Sequence[int],
    layers: Union[int, Sequence[int]],
    *,
    stem_kernel_size: int = 3,
    spatial_padding: int = 96,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Downsampling path of the SPFormer SpConv U-Net.

Embeds the input with a submanifold stem, then runs one SPFormerUNetEncoderBlock per level (blocks); every block but the first downsamples by stride 2 before its residual blocks. The output of every level but the deepest is returned as a skip connection for the decoder.

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • channels (Sequence[int]) –

    Per-level channel widths, deepest level last.

  • layers (Union[int, Sequence[int]]) –

    Number of residual blocks per level; an int is broadcast to every level.

  • stem_kernel_size (int, default: 3 ) –

    Kernel size of the submanifold stem convolution.

  • spatial_padding (int, default: 96 ) –

    Padding (in voxels) added to the inferred spatial shape.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation passed to create_act.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization passed to create_norm.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the normalization.

Shape
  • Input: packed features \((N, \text{in\_channels})\), grid coordinates, batch.
  • Output: SparseConvTensor with channels[-1] channels (bottleneck).

SPFormerUNetDecoder 

SPFormerUNetDecoder(
    channels: Sequence[int],
    layers: Union[int, Sequence[int]],
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Upsampling path of the SPFormer SpConv U-Net.

Runs one SPFormerUNetDecoderBlock per upsampling level (blocks), deepest first. Each block upsamples the deeper feature, concatenates the matching encoder skip, and fuses them back to channels[i] channels. A final normalization and activation (output_layer) is applied to the full-resolution features.

Parameters:

  • channels (Sequence[int]) –

    Per-level channel widths, deepest level last (same as the encoder).

  • layers (Union[int, Sequence[int]]) –

    Number of residual blocks per level; an int is broadcast to every level.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation passed to create_act.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization passed to create_norm.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the normalization.

Shape
  • Input: SparseConvTensor with channels[-1] channels (bottleneck).
  • Output: SparseConvTensor with channels[0] channels.

SPFormerUNetSegmentation 

SPFormerUNetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    channels: Sequence[int] = (32, 64, 96, 128, 160),
    layers: Union[int, Sequence[int]] = 2,
    stem_kernel_size: int = 3,
    spatial_padding: int = 96,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: SegmentationModel

SpConv U-Net from SPFormer.

Reference: sunjiahao1999/SPFormer. A symmetric submanifold sparse-convolution U-Net with pre-norm residual blocks. Distinct from SparseUNetSegmentation: the residual blocks are pre-norm and the stem-level blocks run at full resolution before any downsampling.

Set num_classes=0 to drop the classifier: forward then returns the normalized per-voxel features \((N, \text{channels}[0])\).

Parameters:

  • in_channels (int) –

    Number of input feature channels.

  • num_classes (int) –

    Number of output classes; 0 yields an identity head.

  • channels (Sequence[int], default: (32, 64, 96, 128, 160) ) –

    Per-level channel widths, deepest level last.

  • layers (Union[int, Sequence[int]], default: 2 ) –

    Number of residual blocks per level; an int is broadcast to every level.

  • stem_kernel_size (int, default: 3 ) –

    Kernel size of the submanifold stem convolution.

  • spatial_padding (int, default: 96 ) –

    Padding (in voxels) added to the inferred spatial shape.

  • act (Union[str, Callable, None], default: 'relu' ) –

    Activation passed to create_act.

  • act_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the activation.

  • norm (Union[str, Callable, None], default: 'batch_norm' ) –

    Normalization passed to create_norm.

  • norm_kwargs (Optional[Dict[str, Any]], default: None ) –

    Extra keyword arguments for the normalization (e.g. eps, momentum).

Shape
  • Input: packed features \((N, \text{in\_channels})\), grid coordinates, batch.
  • Output: \((N, \text{num\_classes})\) logits, or \((N, \text{channels}[0])\) features when num_classes=0.

Methods:

  • configure_encoder –

    Builds the sparse encoder producing the bottleneck features and the per-stage skips.

  • configure_decoder –

    Builds the sparse decoder upsampling the bottleneck back to full resolution.

Attributes:

  • num_features (int) –

    Channel count \(C\) of the full-resolution decoder features entering the head.

num_features property 

num_features: int

Channel count \(C\) of the full-resolution decoder features entering the head.

configure_encoder 

configure_encoder() -> SPFormerUNetEncoder

Builds the sparse encoder producing the bottleneck features and the per-stage skips.

configure_decoder 

configure_decoder() -> SPFormerUNetDecoder

Builds the sparse decoder upsampling the bottleneck back to full resolution.