Skip to content

voxel_partition

K-pass voxel-partition inferer with per-point scatter-back aggregation.

Classes:

VoxelPartitionInferer 

VoxelPartitionInferer(
    voxel_size: float,
    transform: Optional[
        Callable[[Dict[str, Any]], Dict[str, Any]]
    ] = None,
    sub_batch_size: int = 1,
    softmax: bool = False,
    reduce: Literal["mean", "sum"] = "mean",
    pos_key: str = POS,
    batch_key: str = BATCH,
    seed: Optional[int] = None,
)

Bases: Inferer

\(K\)-pass voxel-partition inferer with scatter-back aggregation.

Partitions each batch element's points into FNV voxel buckets at voxel_size and runs the predictor on \(K = \max_v c_v\) sub-clouds per element, where sub-cloud \(i\) picks the \((i \bmod c_v)\)-th point of every bucket. Per-sub-cloud logits are scatter-summed to original-point indices and divided by per-point participation counts; each point is picked \(\lfloor K / c_v \rfloor\) or \(\lfloor K / c_v \rfloor + 1\) times across the \(K\) passes, so every original point gets at least one prediction.

For test-time augmentation, wrap in TTAInferer: each TTA pass triggers a fresh \(K\)-pass voxel partition under that augmentation.

Predictions are scatter-summed in float64 for stable averaging across passes; the returned tensor is cast back to the predictor's output dtype. An empty scene (\(N = 0\)) returns a \((0, 0)\) tensor: the predictor is never called, so the channel count cannot be inferred.

Parameters:

  • voxel_size (float) –

    Side length of the FNV voxel partition (in the units of pos).

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

    Optional per-sub-cloud callable applied after slicing each sub-cloud out of data. Typical use: the model's registered preprocessing transform. The transform must preserve the sub-cloud's row count; row-altering pipelines (pad, voxelize, ...) are only supported by SlidingWindowInferer via inverse_key.

  • sub_batch_size (int, default: 1 ) –

    Number of sub-clouds packed into one predictor call via collate. >1 amortises FPS / radius costs on the GPU.

  • softmax (bool, default: False ) –

    If True, softmax each predictor output before scatter-summing.

  • reduce (Literal['mean', 'sum'], default: 'mean' ) –

    "mean" divides each point's accumulated predictions by the number of sub-clouds it appeared in; "sum" returns the plain sum. The argmax is the same within one call, but under TTAInferer the counts differ between views (each view is partitioned on its own augmented positions), so "sum" reproduces the reference protocols that add un-normalized probabilities over views and fragments.

  • pos_key (str, default: POS ) –

    Dict key for the position tensor.

  • batch_key (str, default: BATCH ) –

    Dict key for the per-point batch index.

  • seed (Optional[int], default: None ) –

    Optional RNG seed for the per-pass index shuffle.

Example 
from torch_pointcloud.inferers import VoxelPartitionInferer

inferer = VoxelPartitionInferer(voxel_size=0.04, sub_batch_size=4, transform=model.transforms)
logits = inferer(room, predictor=lambda d: model(d["x"], d["pos"], d["batch"]))