tta

Bases: Inferer

Test-time augmentation inferer.

Runs the wrapped base inferer once per augmentation pass and aggregates the per-point predictions across passes.

Three modes are supported:

• Single callable: re-sampled independently each pass. Use for random augmentations such as uniformly random rotation. Requires num_passes.
• Sequence of callables: each element is applied to exactly one pass in order. Use for a fixed view set (e.g. 8 evenly-spaced rotations). num_passes is inferred from the sequence length.
• None: the input is left as-is and the base inferer is simply run num_passes times, voting over its own randomness (random sub-batches, random padding).

Parameters:

• base (Inferer) –

  Underlying Inferer invoked once per pass. Any concrete inferer works: SimpleInferer(), SlidingWindowInferer(...), KNNWindowInferer(...).

• transforms (Union[TransformFn, Sequence[TransformFn], None], default: None ) –

  Single callable (re-sampled each pass), a sequence of callables for fixed views, or None for un-augmented passes. Any Dict[str, Any] -> Dict[str, Any] callable works, including Compose. When a sequence is given, num_passes is ignored.

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

  Number of TTA passes when transforms is a single callable or None. Must be \(\geq 1\).

• include_identity (bool, default: False ) –

  If True, run one extra pass on the un-augmented input before the augmented passes (the "clean + N random views" voting protocol), so the total pass count is num_passes + 1.

• aggregate (AggregateMode, default: 'mean' ) –

  How per-point predictions are combined across passes. "mean" averages per-pass outputs. "ema" maintains an exponential moving average of them.

• ema_smoothing (float, default: 0.95 ) –

  EMA factor \(\alpha \in [0, 1)\) used when aggregate="ema".

• softmax (bool, default: False ) –

  If True, softmax each pass's output before aggregating. Leave False when the base inferer already returns probabilities (e.g. SlidingWindowInferer(softmax=True)).

• pos_key (str, default: POS ) –

  Dict key for the position tensor (used for the empty-output fallback).

• progress (bool, default: False ) –

  If True, show a tqdm progress bar over the passes.

import torch

import torch_pointcloud as tp
from torch_pointcloud.inferers import SlidingWindowInferer, TTAInferer
from torch_pointcloud.transforms import Compose, RandomFlip, RandomRotate

model = tp.create_model("pointnet.s3dis-area5").eval()
data = {
    "pos": torch.rand(4096, 3) * 10,
    "x": torch.rand(4096, 9),
    "batch": torch.zeros(4096, dtype=torch.long),
}

base = SlidingWindowInferer(block_size=6.0)
aug = Compose([
    RandomRotate(keys="pos", angle_range=(-180.0, 180.0), axis=2, p=1.0),
    RandomFlip(keys="pos", axes=[0, 1], p=0.5),
])
inferer = TTAInferer(base=base, transforms=aug, num_passes=4, aggregate="mean")
probs = inferer(data, predictor=lambda d: model(d["x"], d["pos"], d["batch"]))

from torch_pointcloud.inferers import SlidingWindowInferer, TTAInferer
from torch_pointcloud.transforms import Compose, RandomRotate

base = SlidingWindowInferer(block_size=6.0)
views = [Compose([RandomRotate(keys="pos", angle_range=(a, a), axis=2, p=1.0)])
         for a in (0.0, 45.0, 90.0, 135.0, 180.0, 225.0, 270.0, 315.0)]
inferer = TTAInferer(base=base, transforms=views, aggregate="mean")