torch-pointcloud 0.0.2__tar.gz → 0.0.3__tar.gz

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/.gitignore

@@ -164,6 +164,7 @@ cython_debug/
 weights/**
 data/**
 docs/api/**
+docs/assets/data/
 docs/examples/**
 docs/scripts/**
 !docs/examples/*.ipynb

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: torch-pointcloud
-Version: 0.0.2
+Version: 0.0.3
 Summary: PyTorch Point Cloud models, scripts, pretrained weights and more
 Project-URL: Homepage, https://github.com/arthurdjn/pytorch-pointcloud
 Project-URL: Repository, https://github.com/arthurdjn/pytorch-pointcloud
@@ -158,7 +158,7 @@ Every checkpoint ships the transform that turns a raw point cloud into what the
 # Requires torch-scatter, torch-cluster, spconv
 model, info = tp.create_model("ptv3-base.scannet20.pointcept", task="segmentation", pretrained=True, return_info=True)
 info["transform"]  # the preprocessing pipeline of that checkpoint
-info["weights"]["metrics"]  # {"mIoU": 76.29}
+info["weights"]["metrics"]  # {"mIoU": 77.40}
 
 tp.list_models("pointnext*")  # every registered PointNeXt config
 tp.list_models(task="detection", pretrained=True)  # all detection checkpoints

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/README.md

@@ -115,7 +115,7 @@ Every checkpoint ships the transform that turns a raw point cloud into what the
 # Requires torch-scatter, torch-cluster, spconv
 model, info = tp.create_model("ptv3-base.scannet20.pointcept", task="segmentation", pretrained=True, return_info=True)
 info["transform"]  # the preprocessing pipeline of that checkpoint
-info["weights"]["metrics"]  # {"mIoU": 76.29}
+info["weights"]["metrics"]  # {"mIoU": 77.40}
 
 tp.list_models("pointnext*")  # every registered PointNeXt config
 tp.list_models(task="detection", pretrained=True)  # all detection checkpoints

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/pyproject.toml

@@ -16,7 +16,7 @@ include = [
 
 [project]
 name = "torch-pointcloud"
-version = "0.0.2"
+version = "0.0.3"
 description = "PyTorch Point Cloud models, scripts, pretrained weights and more"
 readme = "README.md"
 license = "Apache-2.0"

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/src/torch_pointcloud/inferers/sliding_window.py

@@ -27,7 +27,7 @@ import torch
 from torch import Tensor
 from tqdm import tqdm
 
-from torch_pointcloud.utils.data import DataKeys
+from torch_pointcloud.utils.data import DataKeys, collate
 
 from ._utils import check_batch_alignment, gaussian_weights, index_select_dict, split_chunks
 from .inferer import Inferer
@@ -194,6 +194,7 @@ def sliding_window_inference(
     mode: WindowMode = "constant",
     sigma_scale: float = 0.125,
     roi_num_points: Optional[int] = None,
+    sw_batch_size: int = 1,
     softmax: bool = True,
     aggregate: AggregateMode = "mean",
     transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
@@ -239,13 +240,18 @@ def sliding_window_inference(
             $D_\text{tiled}$ the number of tiled axes.
         sigma_scale: Gaussian sigma scale factor. Only used when
             `mode="gaussian"`.
-        roi_num_points: Optional cap on points per `predictor` call. Blocks
+        roi_num_points: Optional cap on points per batch element. Blocks
             exceeding this are split into random sub-batches and every point is
             predicted exactly once per block pass. `None` passes the whole block
-            in one call. To enforce a fixed-N predictor input, pair the inferer
+            in one element. At the default `sw_batch_size=1` an element is a call,
+            so this also caps the points per `predictor` call. To enforce a fixed-N predictor input, pair the inferer
             with a `DivisiblePad`-style `transform` that pads each block to a
             multiple of `roi_num_points` and writes its source-to-padded index
             map under `inverse_key`.
+        sw_batch_size: Number of batch elements packed into a single `predictor` call. Blocks are
+            independent, so they are collated along the batch index and their predictions are
+            accumulated in block order: only the number of calls changes. Small blocks are
+            dominated by per-call overhead, where a larger value is several times faster.
         softmax: If `True`, softmax each block's logits before accumulating.
             Use `True` when averaging predictions across multiple blocks or
             TTA passes. Set `False` to accumulate raw logits. `"max"` and
@@ -307,6 +313,8 @@ def sliding_window_inference(
         raise ValueError(f"`aggregate` must be 'mean', 'max' or 'vote', got {aggregate!r}.")
     if roi_num_points is not None and roi_num_points < 1:
         raise ValueError(f"`roi_num_points` must be >= 1 or None, got {roi_num_points}.")
+    if sw_batch_size < 1:
+        raise ValueError(f"`sw_batch_size` must be >= 1, got {sw_batch_size}.")
     if padding < 0.0:
         raise ValueError(f"`padding` must be >= 0, got {padding}.")
 
@@ -343,59 +351,97 @@ def sliding_window_inference(
         weights_b = torch.zeros(n_b, device=device, dtype=torch.float32)
         confidence_b = torch.full((n_b,), -1.0, device=device, dtype=torch.float32)
 
-        for point_ids, w, bbox in tqdm(
-            zip(point_groups, weight_groups, bbox_groups),
-            total=len(point_groups),
+        blocks = list(zip(point_groups, weight_groups, bbox_groups))
+        for start in tqdm(
+            range(0, len(blocks), sw_batch_size),
+            total=math.ceil(len(blocks) / sw_batch_size),
             desc=f"batch {int(b)}",
             leave=False,
             disable=not progress,
         ):
-            n_block = int(point_ids.numel())
-            window = index_select_dict(data_b, point_ids, n_b)
-            window[batch_key] = torch.zeros(n_block, device=device, dtype=torch.long)
-            window[block_bbox_key] = bbox
-            if inverse_key is not None:
-                window.pop(inverse_key, None)
-            if transform is not None:
-                window = transform(window)
-
-            n_window = int(window[pos_key].size(0))
-            inverse_map = window.pop(inverse_key, None) if inverse_key is not None else None
-
-            chunks = split_chunks(n_window, roi_num_points, rng)
-
-            window_preds: Optional[Tensor] = None
-            for chunk_local in chunks:
-                sub_window = index_select_dict(window, chunk_local, n_window)
-                sub_window[batch_key] = torch.zeros(chunk_local.numel(), device=device, dtype=torch.long)
-                logits = predictor(sub_window)
-                if window_preds is None:
-                    window_preds = torch.zeros(n_window, int(logits.size(-1)), device=device, dtype=torch.float32)
-                preds = torch.softmax(logits, dim=-1) if softmax or aggregate != "mean" else logits
-                window_preds[chunk_local] = preds.to(window_preds)
-
-            if window_preds is None:
+            windows: List[Dict[str, Any]] = []
+            inverse_maps: List[Optional[Tensor]] = []
+            samples: List[Dict[str, Any]] = []
+            owners: List[Tuple[int, Tensor]] = []
+
+            for point_ids, _, bbox in blocks[start : start + sw_batch_size]:
+                n_block = int(point_ids.numel())
+                window = index_select_dict(data_b, point_ids, n_b)
+                window[batch_key] = torch.zeros(n_block, device=device, dtype=torch.long)
+                window[block_bbox_key] = bbox
+                if inverse_key is not None:
+                    window.pop(inverse_key, None)
+                if transform is not None:
+                    window = transform(window)
+
+                n_window = int(window[pos_key].size(0))
+                inverse_maps.append(window.pop(inverse_key, None) if inverse_key is not None else None)
+                windows.append(window)
+
+                for chunk_local in split_chunks(n_window, roi_num_points, rng):
+                    sub_window = index_select_dict(window, chunk_local, n_window)
+                    samples.append(sub_window)
+                    owners.append((len(windows) - 1, chunk_local))
+
+            if not samples:
                 continue
 
-            preds_at_block = window_preds if inverse_map is None else window_preds[inverse_map]
-            if scores_b is None:
-                scores_b = torch.zeros(n_b, int(window_preds.size(-1)), device=device, dtype=torch.float32)
-
-            if aggregate == "max":
-                confidence, _ = preds_at_block.max(dim=-1)
-                better = confidence > confidence_b[point_ids]
-                scores_b[point_ids[better]] = preds_at_block[better]
-                confidence_b[point_ids[better]] = confidence[better]
-                weights_b[point_ids] = 1.0
-                continue
-
-            if aggregate == "vote":
-                preds_at_block = torch.nn.functional.one_hot(
-                    preds_at_block.argmax(dim=-1), num_classes=int(preds_at_block.size(-1))
-                ).to(scores_b.dtype)
-
-            scores_b.index_add_(0, point_ids, preds_at_block * w.unsqueeze(-1))
-            weights_b.index_add_(0, point_ids, w)
+            parts: List[Tensor] = []
+            for pack_start in range(0, len(samples), sw_batch_size):
+                packed = collate(
+                    samples[pack_start : pack_start + sw_batch_size],
+                    batch_from=pos_key,
+                    batch_key=batch_key,
+                )
+                logits = predictor(packed)
+                parts.append(torch.softmax(logits, dim=-1) if softmax or aggregate != "mean" else logits)
+            preds = torch.cat(parts)
+
+            window_preds: List[Optional[Tensor]] = [None] * len(windows)
+            offset = 0
+            for (index, chunk_local), sample in zip(owners, samples):
+                count = int(sample[pos_key].size(0))
+                chunk_preds = preds[offset : offset + count]
+                offset += count
+                preds_at_window = window_preds[index]
+                if preds_at_window is None:
+                    n_window = int(windows[index][pos_key].size(0))
+                    preds_at_window = torch.zeros(
+                        n_window,
+                        int(chunk_preds.size(-1)),
+                        device=device,
+                        dtype=torch.float32,
+                    )
+                    window_preds[index] = preds_at_window
+                preds_at_window[chunk_local] = chunk_preds.to(preds_at_window)
+
+            for (point_ids, w, _), block_preds, inverse_map in zip(
+                blocks[start : start + sw_batch_size],
+                window_preds,
+                inverse_maps,
+            ):
+                if block_preds is None:
+                    continue
+
+                preds_at_block = block_preds if inverse_map is None else block_preds[inverse_map]
+                if scores_b is None:
+                    scores_b = torch.zeros(n_b, int(block_preds.size(-1)), device=device, dtype=torch.float32)
+
+                if aggregate == "max":
+                    confidence, _ = preds_at_block.max(dim=-1)
+                    better = confidence > confidence_b[point_ids]
+                    scores_b[point_ids[better]] = preds_at_block[better]
+                    confidence_b[point_ids[better]] = confidence[better]
+                    weights_b[point_ids] = 1.0
+                    continue
+
+                if aggregate == "vote":
+                    preds_at_block = torch.nn.functional.one_hot(
+                        preds_at_block.argmax(dim=-1), num_classes=int(preds_at_block.size(-1))
+                    ).to(scores_b.dtype)
+
+                scores_b.index_add_(0, point_ids, preds_at_block * w.unsqueeze(-1))
+                weights_b.index_add_(0, point_ids, w)
 
         if scores_b is None:
             continue
@@ -442,6 +488,7 @@ class SlidingWindowInferer(Inferer):
         mode: WindowMode = "constant",
         sigma_scale: float = 0.125,
         roi_num_points: Optional[int] = None,
+        sw_batch_size: int = 1,
         softmax: bool = True,
         aggregate: AggregateMode = "mean",
         transform: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
@@ -459,6 +506,7 @@ class SlidingWindowInferer(Inferer):
         self.mode = mode
         self.sigma_scale = sigma_scale
         self.roi_num_points = roi_num_points
+        self.sw_batch_size = sw_batch_size
         self.softmax = softmax
         self.aggregate = aggregate
         self.transform = transform
@@ -484,6 +532,7 @@ class SlidingWindowInferer(Inferer):
             mode=self.mode,
             sigma_scale=self.sigma_scale,
             roi_num_points=self.roi_num_points,
+            sw_batch_size=self.sw_batch_size,
             softmax=self.softmax,
             aggregate=self.aggregate,
             transform=self.transform,

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/src/torch_pointcloud/layers/serialized_attention.py

@@ -25,11 +25,25 @@ class RelativePositionalEncoding(nn.Module):
     r"""Table-based relative position bias added to the attention logits of a serialized patch.
 
     Each axis of the clamped relative grid coordinates indexes its own slice of a shared table, and the three
-    per-head biases are summed.
+    per-head biases are summed. Clamping leaves $\text{rpe\_num}^3$ distinct offsets, so the three rows are
+    summed once per *combination* into a lookup table and the bias is a single gather. Summing them per point
+    pair instead would hold $(\text{pairs}, 3)$ offsets and $(\text{pairs}, 3, H)$ gathered rows, which on a
+    large cloud is the peak allocation of the whole model.
 
     Args:
         patch_size: Number of points attending to each other, setting the clamping boundary.
         num_heads: Number of attention heads.
+
+    Shape:
+        - Input: $(P, K, 3)$ integer grid coordinates, $K$ points per patch.
+        - Output: $(P, H, K, K)$ bias, aligned with the attention logits.
+
+    Example:
+        ```python
+        rpe = RelativePositionalEncoding(patch_size=8, num_heads=4)
+        bias = rpe(torch.randint(0, 16, (2, 8, 3)))
+        bias.shape  # (2, 4, 8, 8)
+        ```
     """
 
     def __init__(self, patch_size: int, num_heads: int) -> None:
@@ -41,18 +55,24 @@ class RelativePositionalEncoding(nn.Module):
         self.rpe_table = nn.Parameter(torch.zeros(3 * self.rpe_num, num_heads))
         nn.init.trunc_normal_(self.rpe_table, std=0.02)
 
-    def forward(self, relative_pos: Tensor) -> Tensor:
-        clamped_pos = relative_pos.clamp(-self.coords_boundary, self.coords_boundary)
-        positive_indices = clamped_pos + self.coords_boundary
-        dim_strides = torch.arange(3, device=relative_pos.device) * self.rpe_num
+    def _relative_index(self, coord: Tensor) -> Tensor:
+        r"""Clamped relative offset of every point pair along one axis. $(P, K) \to (P, K, K)$."""
+        relative = coord.unsqueeze(2) - coord.unsqueeze(1)
+        return relative.clamp(-self.coords_boundary, self.coords_boundary) + self.coords_boundary
 
-        idx = positive_indices + dim_strides
+    def forward(self, pos_grid: Tensor) -> Tensor:
+        table = self.rpe_table.view(3, self.rpe_num, self.num_heads)
+        combined = table[0].view(-1, 1, 1, self.num_heads) + table[1].view(1, -1, 1, self.num_heads)
+        combined = combined + table[2].view(1, 1, -1, self.num_heads)
 
-        encodings = self.rpe_table.index_select(0, idx.reshape(-1))
-        encodings = encodings.view(idx.shape + (-1,)).sum(3)
-        encodings = encodings.permute(0, 3, 1, 2)
+        x, y, z = pos_grid.unbind(dim=-1)
+        idx = self._relative_index(x) * self.rpe_num + self._relative_index(y)
+        idx = idx * self.rpe_num + self._relative_index(z)
 
-        return encodings
+        encodings = combined.view(-1, self.num_heads).index_select(0, idx.reshape(-1))
+        encodings = encodings.view(idx.shape + (self.num_heads,))
+
+        return encodings.permute(0, 3, 1, 2)
 
     def extra_repr(self) -> str:
         return f"patch_size={self.patch_size}, num_heads={self.num_heads}"
@@ -124,6 +144,7 @@ class SerializedAttention(nn.Module):
                 raise ValueError("Upcasting attention is not supported with Flash Attention.")
             if upcast_softmax:
                 raise ValueError("Upcasting softmax is not supported with Flash Attention.")
+
         self.channels = channels
         self.num_heads = num_heads
         self.patch_size = patch_size
@@ -155,8 +176,13 @@ class SerializedAttention(nn.Module):
         pad_mode: PadMode = "above" if self.use_flash_attn else "all"
 
         padded_indices, unpadded_indices, padded_batch = divisible_pad(
-            batch, patch_size, mode=pad_mode, pad_fill="replicate", return_inverse=True
+            batch,
+            patch_size,
+            mode=pad_mode,
+            pad_fill="replicate",
+            return_inverse=True,
         )
+
         order = serialized_order[padded_indices] if serialized_order is not None else padded_indices
         inverse = unpadded_indices[serialized_inverse] if serialized_inverse is not None else unpadded_indices
         qkv = self.qkv(x)[order]
@@ -171,9 +197,11 @@ class SerializedAttention(nn.Module):
             if self.upcast_attn:
                 q = q.float()
                 k = k.float()
+
             attn = (q * self.scale) @ k.transpose(-2, -1)
             if self.upcast_softmax:
                 attn = attn.float()
+
             attn = attn.softmax(dim=-1)
             attn = F.dropout(attn, p=self.attn_drop, training=self.training).to(qkv.dtype)
             feat = (attn @ v).transpose(1, 2).reshape(-1, C)
@@ -243,8 +271,13 @@ class SerializedAttentionRPE(nn.Module):
         patch_size = self.patch_size
 
         padded_indices, unpadded_indices, _ = divisible_pad(
-            batch, patch_size, mode="all", pad_fill="replicate", return_inverse=True
+            batch,
+            patch_size,
+            mode="all",
+            pad_fill="replicate",
+            return_inverse=True,
         )
+
         order = serialized_order[padded_indices] if serialized_order is not None else padded_indices
         inverse = unpadded_indices[serialized_inverse] if serialized_inverse is not None else unpadded_indices
         qkv = self.qkv(x)[order]
@@ -255,14 +288,12 @@ class SerializedAttentionRPE(nn.Module):
         if self.upcast_attn:
             q = q.float()
             k = k.float()
-        attn = (q * self.scale) @ k.transpose(-2, -1)
-
-        pos_grid_ordered = pos_grid_ordered.reshape(-1, K, 3)
-        relative_pos = pos_grid_ordered.unsqueeze(2) - pos_grid_ordered.unsqueeze(1)
-        attn = attn + self.rpe(relative_pos)
 
+        attn = (q * self.scale) @ k.transpose(-2, -1)
+        attn = attn + self.rpe(pos_grid_ordered.reshape(-1, K, 3))
         if self.upcast_softmax:
             attn = attn.float()
+
         attn = attn.softmax(dim=-1)
         attn = F.dropout(attn, p=self.attn_drop, training=self.training).to(qkv.dtype)
         feat = (attn @ v).transpose(1, 2).reshape(-1, C)
@@ -313,6 +344,7 @@ class SerializedAttentionRoPE(nn.Module):
                 raise ValueError("Upcasting attention is not supported with Flash Attention.")
             if upcast_softmax:
                 raise ValueError("Upcasting softmax is not supported with Flash Attention.")
+
         self.channels = channels
         self.num_heads = num_heads
         self.patch_size = patch_size
@@ -344,8 +376,13 @@ class SerializedAttentionRoPE(nn.Module):
         pad_mode: PadMode = "above" if self.use_flash_attn else "all"
 
         padded_indices, unpadded_indices, padded_batch = divisible_pad(
-            batch, patch_size, mode=pad_mode, pad_fill="replicate", return_inverse=True
+            batch,
+            patch_size,
+            mode=pad_mode,
+            pad_fill="replicate",
+            return_inverse=True,
         )
+
         order = serialized_order[padded_indices] if serialized_order is not None else padded_indices
         inverse = unpadded_indices[serialized_inverse] if serialized_inverse is not None else unpadded_indices
         qkv = self.qkv(x)[order]
@@ -366,9 +403,11 @@ class SerializedAttentionRoPE(nn.Module):
             if self.upcast_attn:
                 q = q.float()
                 k = k.float()
+
             attn = (q * self.scale) @ k.transpose(-2, -1)
             if self.upcast_softmax:
                 attn = attn.float()
+
             attn = attn.softmax(dim=-1)
             attn = F.dropout(attn, p=self.attn_drop, training=self.training).to(qkv.dtype)
             feat = (attn @ v).transpose(1, 2).reshape(-1, C)

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/src/torch_pointcloud/models/_base.py

@@ -87,8 +87,11 @@ class ClassificationModel(nn.Module, metaclass=ABCMeta):
     def forward_features(self, *args: Any, **kwargs: Any) -> Any:
         r"""Encode a packed point cloud into features for `forward_head`.
 
-        Canonical signature: `forward_features(x, pos, batch)`, returning the encoded features (often with
-        the downsampled `pos` / `batch`).
+        Canonical signature: `forward_features(x, pos, batch, return_intermediates=False)`, returning the encoded
+        features (often with the downsampled `pos` / `batch`). With `return_intermediates=True`, the finer encoder
+        stages are appended as the last element of the return: a `List[FeaturesDict]` in fine-to-coarse order, each
+        entry holding `x`, `batch` and the stage's `pos` or `pos_grid`. Models add the keys their decoder needs or use
+        their own container when a stage is not a packed point set (dense tokens, voxel-point pairs).
         """
 
     @abstractmethod
@@ -157,8 +160,11 @@ class SegmentationModel(nn.Module, metaclass=ABCMeta):
     def forward_features(self, *args: Any, **kwargs: Any) -> Any:
         r"""Encode a packed point cloud, keeping what `forward_decoder` needs for the skip connections.
 
-        Canonical signature: `forward_features(x, pos, batch)`, returning the encoder output together with
-        the per-stage intermediates.
+        Canonical signature: `forward_features(x, pos, batch, return_intermediates=False)`, returning the encoder
+        output. With `return_intermediates=True`, the finer encoder stages are appended as the last element of the
+        return: a `List[FeaturesDict]` in fine-to-coarse order, each entry holding `x`, `batch` and the stage's `pos`
+        or `pos_grid`. Models add the keys their decoder needs (e.g. `pooling_inverse`) or use their own container
+        when a stage is not a packed point set (dense tokens, voxel-point pairs).
         """
 
     def forward_decoder(self, *args: Any, **kwargs: Any) -> Any:

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/src/torch_pointcloud/models/concerto.py

@@ -14,7 +14,11 @@ import torch_pointcloud.transforms as T
 from torch_pointcloud.datasets.scannet import SCANNET20_CLASSES
 from torch_pointcloud.models._base import SegmentationModel
 from torch_pointcloud.models._registry import WeightsDict, register_model
-from torch_pointcloud.models.point_transformer_v3 import AttentionKind, PointTransformerV3Encoder
+from torch_pointcloud.models.point_transformer_v3 import (
+    AttentionKind,
+    PointTransformerV3Encoder,
+    SerializedFeaturesDict,
+)
 from torch_pointcloud.utils.data import DataKeys
 from torch_pointcloud.utils.serialization import SerializationOrder
 from torch_pointcloud.utils.types import OptTensor
@@ -154,7 +158,7 @@ class ConcertoSegmentation(SegmentationModel):
         pos_grid: Tensor,
         batch: Tensor,
         return_intermediates: Literal[True],
-    ) -> Tuple[Tensor, Tensor, Tensor, List[Dict[str, Tensor]]]: ...
+    ) -> Tuple[Tensor, Tensor, Tensor, List[SerializedFeaturesDict]]: ...
 
     @overload
     def forward_features(
@@ -176,10 +180,10 @@ class ConcertoSegmentation(SegmentationModel):
             return self.encoder.forward(x, pos_grid, batch, return_intermediates=True)
         return self.encoder.forward(x, pos_grid, batch, return_intermediates=False)
 
-    def forward_decoder(self, x: Tensor, intermediates: List[Dict[str, Tensor]]) -> Tuple[Tensor, Tensor, Tensor]:
+    def forward_decoder(self, x: Tensor, intermediates: List[SerializedFeaturesDict]) -> Tuple[Tensor, Tensor, Tensor]:
         pos_grid = batch = None
         for intermediate in reversed(intermediates):
-            inverse = intermediate["inverse"]
+            inverse = intermediate["pooling_inverse"]
             x = torch.cat([intermediate["x"], x[inverse]], dim=-1)
             pos_grid = intermediate["pos_grid"]
             batch = intermediate["batch"]

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/src/torch_pointcloud/models/dgcnn.py

@@ -3,7 +3,7 @@
 {{ paper("1801.07829") }}
 """
 
-from typing import Any, Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -26,13 +26,6 @@ from ._base import ClassificationModel, SegmentationModel
 from ._registry import WeightsDict, register_model
 
 
-class DGCNNIntermediate(NamedTuple):
-    """Per-block encoder features and their batch index."""
-
-    x: Tensor
-    batch: Tensor
-
-
 class DGCNNEncoderBlock(nn.Module):
     """EdgeConv block: rebuilds a $k$-nearest-neighbor graph in feature space and runs an MLP over the edges.
 

{torch_pointcloud-0.0.2 → torch_pointcloud-0.0.3}/src/torch_pointcloud/models/kpconv.py

@@ -27,7 +27,7 @@ from torch_pointcloud.utils.conversion import ensure_tuple, ensure_tuple_size
 from torch_pointcloud.utils.data import DataKeys
 from torch_pointcloud.utils.geometry import rodrigues_rotation_matrix, spherical_points_gradient, spherical_points_lloyd
 from torch_pointcloud.utils.imports import _TORCH_CLUSTER_GITHUB_URL, _TORCH_SCATTER_GITHUB_URL, optional_import
-from torch_pointcloud.utils.types import OptTensor
+from torch_pointcloud.utils.types import OptTensor, PooledFeaturesDict
 
 from ._base import ClassificationModel, SegmentationModel
 from ._registry import WeightsDict, register_model
@@ -770,7 +770,7 @@ class KPFCNNEncoder(nn.Module):
         pos: Tensor,
         batch: Tensor,
         return_intermediates: Literal[True],
-    ) -> Tuple[Tensor, Tensor, Tensor, List[Dict[str, Tensor]]]: ...
+    ) -> Tuple[Tensor, Tensor, Tensor, List[PooledFeaturesDict]]: ...
 
     @overload
     def forward(
@@ -788,12 +788,12 @@ class KPFCNNEncoder(nn.Module):
         batch: Tensor,
         return_intermediates: bool = False,
     ) -> Any:
-        intermediates = []
+        intermediates: List[PooledFeaturesDict] = []
         for i, block in enumerate(self.blocks):
-            intermediate = {"x": x, "pos": pos, "batch": batch}
+            intermediate: PooledFeaturesDict = {"x": x, "pos": pos, "batch": batch}
             x, pos, batch, inv = block(x, pos, batch, return_inverse=True)
 
-            if i > 0:
+            if return_intermediates and i > 0:
                 intermediate["pooling_inverse"] = inv
                 intermediates.append(intermediate)
 
@@ -994,7 +994,7 @@ class KPFCNNClassification(ClassificationModel):
         pos: Tensor,
         batch: Tensor,
         return_intermediates: Literal[True],
-    ) -> Tuple[Tensor, Tensor, Tensor, List[Dict[str, Tensor]]]: ...
+    ) -> Tuple[Tensor, Tensor, Tensor, List[PooledFeaturesDict]]: ...
 
     @overload
     def forward_features(
@@ -1291,7 +1291,7 @@ class KPFCNNSegmentation(SegmentationModel):
         pos: Tensor,
         batch: Tensor,
         return_intermediates: Literal[True],
-    ) -> Tuple[Tensor, Tensor, Tensor, List[Dict[str, Tensor]]]: ...
+    ) -> Tuple[Tensor, Tensor, Tensor, List[PooledFeaturesDict]]: ...
 
     @overload
     def forward_features(
@@ -1331,20 +1331,19 @@ class KPFCNNSegmentation(SegmentationModel):
             else:
                 x = self.stem(x)
 
-        intermediates = [{"x": x, "pos": pos, "batch": batch}]
-        x, pos, batch, encoder_intermediates = self.encoder(x, pos, batch, return_intermediates=True)
-        intermediates.extend(encoder_intermediates)
+        if not return_intermediates:
+            return self.encoder(x, pos, batch)
 
-        if return_intermediates:
-            return x, pos, batch, intermediates
-        return x, pos, batch
+        skip: PooledFeaturesDict = {"x": x, "pos": pos, "batch": batch}
+        x, pos, batch, intermediates = self.encoder(x, pos, batch, return_intermediates=True)
+        return x, pos, batch, [skip, *intermediates]
 
     def forward_decoder(
         self,
         x: Tensor,
         pos: Tensor,
         batch: Tensor,
-        intermediates: List[Dict[str, Tensor]],
+        intermediates: List[PooledFeaturesDict],
     ) -> Tensor:
         x, _, _ = self.decoder(x, pos, batch, intermediates)
         return x