SPVCNN

SPVCNN classification and segmentation models.

First page of Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution

2007.16100 · July 2020

Classes:

PointTensor –

A SparseTensor subclass that caches per-stride point↔voxel mappings.
BasicBlock –

Sparse 3D convolution followed by normalization and activation.
ResidualBlock –

Residual block of two sparse 3D convolutions.
SPVFusionBlock –

Point-voxel fusion: devoxelizes the sparse features and adds a linear projection of the point branch.
SPVCNNUpsampleBlock –

Transposed sparse convolution that doubles the resolution, then a residual block over the concatenated skip.
SPVCNNEncoderBlock –

One encoder stage: an optional downsampling convolution, depth residual blocks, and an optional fusion block.
SPVCNNDecoderBlock –

One decoder stage: an optional upsampling block, depth residual blocks, and an optional fusion block.
SPVCNNIntermediateDict –

Per-stage encoder features kept for the decoder skip connections.
SPVCNNEncoder –

Stack of SPVCNNEncoderBlock stages, halving the resolution and fusing the point branch at selected stages.
SPVCNNDecoder –

Stack of SPVCNNDecoderBlock stages consuming the encoder intermediates from coarsest to finest resolution.
SPVCNNClassification –

SPVCNN classification model as described in the paper
SPVCNNSegmentation –

SPVCNN segmentation model as described in the paper

Functions:

initial_voxelize –

Aggregate a PointTensor into a SparseTensor of voxel features.
point_to_voxel –

Aggregate point features (z.F) onto the voxel grid of x.
voxel_to_point –

Trilinearly interpolate voxel features (x.F) at point positions (z.C).

PointTensor ¶

PointTensor(
    feats: Tensor,
    coords: Tensor,
    stride: Union[int, Tuple[int, ...]] = 1,
)

Bases: SparseTensor

A SparseTensor subclass that caches per-stride point↔voxel mappings.

The voxelization helpers (initial_voxelize, point_to_voxel, voxel_to_point) expect coordinates in batch-FIRST layout [B, X, Y, Z] (matching torchsparse's SparseTensor.C).

BasicBlock ¶

BasicBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    transposed: bool = False,
    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 3D convolution followed by normalization and activation.

Set transposed=True for an upsampling (inverse) convolution.

ResidualBlock ¶

ResidualBlock(
    in_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.0,
    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

Residual block of two sparse 3D convolutions.

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

SPVFusionBlock ¶

SPVFusionBlock(
    in_channels: int,
    out_channels: 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

Point-voxel fusion: devoxelizes the sparse features and adds a linear projection of the point branch.

The fused point features are voxelized back so that both branches carry the fusion.

SPVCNNUpsampleBlock ¶

SPVCNNUpsampleBlock(
    in_channels: int,
    skip_channels: int,
    out_channels: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    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

Transposed sparse convolution that doubles the resolution, then a residual block over the concatenated skip.

SPVCNNEncoderBlock ¶

SPVCNNEncoderBlock(
    in_channels: int,
    out_channels: int,
    depth: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: Union[float, Sequence[float]] = 0.0,
    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,
    fusion: Optional[Module] = None,
    downsample: Optional[Module] = None,
)

Bases: Module

One encoder stage: an optional downsampling convolution, depth residual blocks, and an optional fusion block.

Parameters:

fusion (Optional[Module], default: None ) –

Point-voxel fusion applied after the residual blocks. Requires point features at forward time.
downsample (Optional[Module], default: None ) –

Strided convolution applied before the residual blocks.

SPVCNNDecoderBlock ¶

SPVCNNDecoderBlock(
    channels: int,
    depth: int,
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    dropout: float = 0.0,
    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,
    fusion: Optional[Module] = None,
    upsample: Optional[Module] = None,
)

Bases: Module

One decoder stage: an optional upsampling block, depth residual blocks, and an optional fusion block.

Parameters:

fusion (Optional[Module], default: None ) –

Point-voxel fusion applied after the residual blocks. Requires point features at forward time.
upsample (Optional[Module], default: None ) –

Upsampling block merging the encoder skip. Requires x_voxels_skip at forward time.

SPVCNNIntermediateDict ¶

Bases: TypedDict

Per-stage encoder features kept for the decoder skip connections.

SPVCNNEncoder ¶

SPVCNNEncoder(
    channels: Sequence[int],
    depths: Sequence[int],
    fusion_stages: Sequence[bool],
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.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

Stack of SPVCNNEncoderBlock stages, halving the resolution and fusing the point branch at selected stages.

SPVCNNDecoder ¶

SPVCNNDecoder(
    depths: Sequence[int],
    channels: Sequence[int],
    skip_channels: Sequence[int],
    fusion_stages: Sequence[bool],
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    dropout: float = 0.0,
    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

Stack of SPVCNNDecoderBlock stages consuming the encoder intermediates from coarsest to finest resolution.

SPVCNNClassification ¶

SPVCNNClassification(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stem_channels: int = 32,
    encoder_channels: Sequence[int] = (32, 64, 128, 256),
    encoder_depths: Sequence[int] = (2, 2, 2, 2),
    encoder_fusion_stages: Sequence[bool] = (
        False,
        False,
        False,
        True,
    ),
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.3,
    global_pool: PoolLike = "max",
    dropout: float = 0.0,
    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: ClassificationModel

SPVCNN classification model as described in the paper Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution by Haotian Tang, Zhijian Liu, Shengyu Zhao, Yujun Lin, Ji Lin, Hanrui Wang, Song Han.

Sparse voxel convolutions carry the coarse context while a parallel point branch keeps the fine geometry, the two being merged at every fusion stage.

Parameters:

in_channels (int) –

Number of input channels. Falls back to spatial_dim when None.
num_classes (int) –

Number of output classes.
spatial_dim (int, default: 3 ) –

Spatial dimensionality of the point clouds.
stem_channels (int, default: 32 ) –

Number of channels of the two stem convolutions.
encoder_channels (Sequence[int], default: (32, 64, 128, 256) ) –

Output channels of each encoder stage.
encoder_depths (Sequence[int], default: (2, 2, 2, 2) ) –

Number of residual blocks in each encoder stage.
encoder_fusion_stages (Sequence[bool], default: (False, False, False, True) ) –

Whether each encoder stage ends with a point-voxel fusion block.
kernel_size (int, default: 3 ) –

Kernel size of the residual convolutions.
stride (int, default: 1 ) –

Stride of the residual convolutions.
dilation (int, default: 1 ) –

Dilation of the residual convolutions.
drop_path (float, default: 0.3 ) –

Maximum stochastic depth rate, ramped linearly across the encoder blocks.
global_pool (PoolLike, default: 'max' ) –

Global pooling used to reduce the point features to one vector per cloud.
dropout (float, default: 0.0 ) –

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

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

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

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

Optional keyword arguments for the normalization factory.

Methods:

configure_stem –

Build the two-convolution sparse stem lifting the input features to stem_channels.
configure_encoder –

Build the SPVCNNEncoder 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_stem ¶

configure_stem() -> Sequential

Build the two-convolution sparse stem lifting the input features to stem_channels.

configure_encoder ¶

configure_encoder() -> SPVCNNEncoder

Build the SPVCNNEncoder backbone.

SPVCNNSegmentation ¶

SPVCNNSegmentation(
    in_channels: int,
    num_classes: int,
    *,
    spatial_dim: int = 3,
    stem_channels: int = 32,
    encoder_channels: Sequence[int],
    encoder_depths: Sequence[int],
    encoder_fusion_stages: Sequence[bool],
    decoder_channels: Sequence[int],
    decoder_depths: Sequence[int],
    decoder_fusion_stages: Sequence[bool],
    kernel_size: int = 3,
    stride: int = 1,
    dilation: int = 1,
    drop_path: float = 0.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: SegmentationModel

SPVCNN segmentation model as described in the paper Searching Efficient 3D Architectures with Sparse Point-Voxel Convolution by Haotian Tang, Zhijian Liu, Shengyu Zhao, Yujun Lin, Ji Lin, Hanrui Wang, Song Han.

A sparse voxel U-Net with skip connections, run alongside a point branch that the fusion stages merge back into the voxel features. Per-point logits are read from the point branch.

Parameters:

in_channels (int) –

Number of input channels. Falls back to spatial_dim when None.
num_classes (int) –

Number of output classes.
spatial_dim (int, default: 3 ) –

Spatial dimensionality of the point clouds.
stem_channels (int, default: 32 ) –

Number of channels of the two stem convolutions.
encoder_channels (Sequence[int]) –

Output channels of each encoder stage.
encoder_depths (Sequence[int]) –

Number of residual blocks in each encoder stage.
encoder_fusion_stages (Sequence[bool]) –

Whether each encoder stage ends with a point-voxel fusion block.
decoder_channels (Sequence[int]) –

Output channels of each decoder stage.
decoder_depths (Sequence[int]) –

Number of residual blocks in each decoder stage.
decoder_fusion_stages (Sequence[bool]) –

Whether each decoder stage ends with a point-voxel fusion block.
kernel_size (int, default: 3 ) –

Kernel size of the residual convolutions.
stride (int, default: 1 ) –

Stride of the residual convolutions.
dilation (int, default: 1 ) –

Dilation of the residual convolutions.
drop_path (float, default: 0.3 ) –

Maximum stochastic depth rate, ramped linearly across the encoder blocks.
act (Union[str, Callable, None], default: 'relu' ) –

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

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

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

Optional keyword arguments for the normalization factory.

Methods:

configure_stem –

Build the two-convolution sparse stem lifting the input features to stem_channels.
configure_encoder –

Build the SPVCNNEncoder backbone.
configure_decoder –

Build the SPVCNNDecoder 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_stem ¶

configure_stem() -> Sequential

Build the two-convolution sparse stem lifting the input features to stem_channels.

configure_encoder ¶

configure_encoder() -> SPVCNNEncoder

Build the SPVCNNEncoder backbone.

configure_decoder ¶

configure_decoder() -> SPVCNNDecoder

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

initial_voxelize ¶

initial_voxelize(
    z: PointTensor,
    init_res: float = 1.0,
    after_res: float = 1.0,
) -> SparseTensor

Aggregate a PointTensor into a SparseTensor of voxel features.

Mutates z.C to the rescaled (voxel-unit) float coordinates so subsequent voxel_to_point calls can stay in voxel space.

point_to_voxel ¶

point_to_voxel(
    x: SparseTensor, z: PointTensor
) -> SparseTensor

Aggregate point features (z.F) onto the voxel grid of x.

voxel_to_point ¶

voxel_to_point(
    x: SparseTensor, z: PointTensor, nearest: bool = False
) -> PointTensor

Trilinearly interpolate voxel features (x.F) at point positions (z.C).

SPVCNN

PointTensor ¶

BasicBlock ¶

ResidualBlock ¶

SPVFusionBlock ¶

SPVCNNUpsampleBlock ¶

SPVCNNEncoderBlock ¶

SPVCNNDecoderBlock ¶

SPVCNNIntermediateDict ¶

SPVCNNEncoder ¶

SPVCNNDecoder ¶

SPVCNNClassification ¶

num_features property ¶

configure_stem ¶

configure_encoder ¶

SPVCNNSegmentation ¶

num_features property ¶

configure_stem ¶

configure_encoder ¶

configure_decoder ¶

initial_voxelize ¶

point_to_voxel ¶

voxel_to_point ¶

num_features `property` ¶

num_features `property` ¶