minuet.nn.convolutions#

Classes

KernelMapCache(ndim, dtype, device[, ...])

KernelMapCache builds kernel map without duplicates.

SparseConv(ndim, in_channels, out_channels, ...)

Applies a sparse convolution over a SparseTensor.

SparseConv3d(in_channels, out_channels, ...)

Applies a 3-D sparse convolution over an input point cloud and produces an output point cloud.
class KernelMapCache(ndim: int, dtype: dtype, device: str | device, layout: str = 'minuet', disable_reordering: bool = False)#

KernelMapCache builds kernel map without duplicates. It works jointly with SparseConv.

class SparseConv(ndim: int, in_channels: int, out_channels: int, kernel_size: int | Tuple[int, ...], stride: int | Tuple[int, ...] = 1, dilation: int | Tuple[int, ...] = 1, bias: bool = False, transposed: bool = False, dtype: dtype | None = None, device: device | None = None)#

Applies a sparse convolution over a SparseTensor.

Each input point cloud in the SparseTensor consists of a set of non-zero points (i.e. the coordinate tensor of the SparseTensor) \(\mathcal{P} = \{\mathbf{p_j}\}\) and its corresponding features \(\{\mathbf{F^\mathcal{P}_j}\}\) (i.e. the feature tensor of the SparseTensor). With the stride \(s\), the sparse convolution will produce an output point cloud, where the set of non-zero points are generated as follows:

\[\mathcal{Q} = \left\{ \left( \left\lfloor \frac{x}{s} \right\rfloor \times s, \left\lfloor \frac{y}{s} \right\rfloor \times s, \left\lfloor \frac{z}{s} \right\rfloor \times s \right) ~\middle|~ (x, y, z) \in \mathcal{P} \right\}\]

The output feature vector \(\mathbf{F^\mathcal{Q}_i}\) of each output coordinate \(\mathbf{q_i}\) is computed as follows:

\[\mathbf{F^\mathcal{Q}_i} = \sum_{\mathbf{p_j} \in \mathcal{P}} \sum_{\boldsymbol{\delta_k} \in \Delta} 𝟙_{\mathbf{p_j} = \mathbf{q_i} + \boldsymbol{\delta_k}} \mathbf{F^\mathcal{P}_j} \mathbf{W_k}\]

The \(\Delta\) denotes the weight offsets, generated by generate_kernel_offsets().

Shape:

If there are \(N\) input points and \(M\) output points, then:

  • Input coordinate tensor: \([N, 3]\)

  • Input feature tensor: \([N, C_\text{in}]\)

  • Output coordinate tensor: \([M, 3]\)

  • Output feature tensor: \([M, C_\text{out}]\)

  • Kernel weight offsets: \([K^3, 3]\)

  • Kernel weights: \([K^3, C_\text{in}, C_\text{out}]\)

  • Bias (if enabled): \([C_\text{out}]\)

Parameters:

  • ndim – the number of coordinate dimensions

  • in_channels – the number of input feature channels \(C_\text{in}\)

  • out_channels – the number of input feature channels \(C_\text{out}\)

  • kernel_size – the kernel size of the sparse convolution \(K\)

  • stride – the stride of the sparse convolution \(s\)

  • dilation – the dilation of the sparse convolution

  • bias – whether enables bias weight in the sparse convolution

  • transposed – whether enables transposed sparse convolution

  • dtype – the data type of the weights

  • device – the device to store the weights

autotune(free_buffers: bool = True)#

A context for autotuning the current SparseConv.

Examples

conv = SparseConv(...)
with conv.autotune():
  for inputs in data_loader:
    _ = conv(inputs)
Parameters:

free_buffers – whether to release all buffers after autotuning for saving GPU memory

reset_parameters()#

Randomly initialize all weights for the SparseConv

set_kernel_map_cache(cache: KernelMapCache)#

Set the kernel map cache for the current SparseConv

Parameters:

cache – the cache to be attached to the current SparseConv

property tunable_config: Dict[str, Any]#

Returns: all tunable configurations for the SparseConv

class SparseConv3d(in_channels: int, out_channels: int, kernel_size: int | Tuple[int, ...], stride: int | Tuple[int, ...] = 1, dilation: int | Tuple[int, ...] = 1, bias: bool = False, dtype: dtype | None = None, device: device | None = None)#

Applies a 3-D sparse convolution over an input point cloud and produces an output point cloud. Please refer to SparseConv for the meanings of each parameter.