SpUNet

SpUNet segmentation model.

First page of 4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks

1904.08755 · April 2019

Classes:

SparseBasicBlock –

Residual block of two submanifold sparse convolutions.
SparseUNetEncoder –

Sparse convolutional stem followed by stages that halve the resolution and run SparseBasicBlock blocks.
SparseUNetDecoder –

Mirror of SparseUNetEncoder: each stage inverse-convolves, concatenates its skip, and runs residual blocks.
SparseUNetSegmentation –

SpUNet segmentation model, a sparse residual U-Net in the spirit of

SparseBasicBlock ¶

SparseBasicBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    bias: bool = False,
    indice_key: Optional[str] = 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: SparseModule

Residual block of two submanifold sparse convolutions.

A pointwise convolution projects the skip connection when the channel count changes.

SparseUNetEncoder ¶

SparseUNetEncoder(
    in_channels: int,
    base_channels: int = 32,
    channels: Sequence[int] = (32, 64, 128, 256),
    layers: Sequence[int] = (2, 3, 4, 6),
    stem_kernel_size: int = 5,
    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

Sparse convolutional stem followed by stages that halve the resolution and run SparseBasicBlock blocks.

The stem output and every stage output but the last are returned as decoder skips.

SparseUNetDecoder ¶

SparseUNetDecoder(
    in_channels: int,
    skip_channels: Sequence[int],
    channels: Sequence[int],
    layers: Sequence[int],
    kernel_size: int = 3,
    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

Mirror of SparseUNetEncoder: each stage inverse-convolves, concatenates its skip, and runs residual blocks.

Stages are built shallowest-first but run deepest-first, so channels and layers are read back-to-front.

SparseUNetSegmentation ¶

SparseUNetSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    base_channels: int = 32,
    channels: Sequence[int] = (
        32,
        64,
        128,
        256,
        256,
        128,
        96,
        96,
    ),
    layers: Sequence[int] = (2, 3, 4, 6, 2, 2, 2, 2),
    stem_kernel_size: int = 5,
    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

SpUNet segmentation model, a sparse residual U-Net in the spirit of 4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks by Christopher Choy, JunYoung Gwak, Silvio Savarese.

Voxelized coordinates enter a sparse convolutional stem, then a symmetric encoder-decoder of submanifold residual blocks with skip connections. The head runs on the full-resolution decoder output, whose rows stay aligned with the input points.

Parameters:

in_channels (int) –

Number of input feature channels. Positions are used as features when x is None.
num_classes (int) –

Number of output classes. \(0\) replaces the head with nn.Identity.
base_channels (int, default: 32 ) –

Number of channels of the stem, also the width of the shallowest skip.
channels (Sequence[int], default: (32, 64, 128, 256, 256, 128, 96, 96) ) –

Feature width of every stage, encoder stages first then decoder stages. Must have even length.
layers (Sequence[int], default: (2, 3, 4, 6, 2, 2, 2, 2) ) –

Number of residual blocks per stage, aligned with channels.
stem_kernel_size (int, default: 5 ) –

Kernel size of the stem convolution.
kernel_size (int, default: 3 ) –

Kernel size of the residual convolutions.
spatial_padding (int, default: 96 ) –

Padding added to the sparse spatial shape, so that voxel indices stay in range.
act (Union[str, Callable, None], default: 'relu' ) –

Activation function.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the activation function.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization function.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Keyword arguments for the normalization function.

Methods:

configure_encoder –

Build the sparse encoder producing the bottleneck features and the per-stage skips.
configure_decoder –

Build 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() -> SparseUNetEncoder

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

configure_decoder ¶

configure_decoder() -> SparseUNetDecoder

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