voxel_partition

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 per-sub-cloud part of the model's preprocessing (centering, grid coordinates, feature stacking). If it changes the row count (pad, voxelize, ...) it must record a source-to-predictor index map under inverse_key so the inferer can gather predictions back to the sub-cloud's points.

• 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.

• aggregate (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.

• inverse_key (Optional[str], default: None ) –

  Dict key under which a row-altering transform records a source-to-predictor long index map of shape \((N_\text{sub},)\) with values in \([0, N_\text{predictor})\). Any scene-level value at this key is dropped before transform runs, and the map is popped before the predictor is called. Leave None when the transform preserves row count.

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

  RNG seed for the per-pass index shuffle. None draws from the global generator, so torch.manual_seed seeds the inferer together with the transforms. An int is offset by the number of calls the instance has made: repeated calls (e.g. TTAInferer views) draw different shuffles, and a fresh instance replays the same sequence.

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"]))