From 359aca94eab008bfd219cbe3e604afc3f05015c5 Mon Sep 17 00:00:00 2001 From: dripston Date: Sun, 20 Sep 2026 00:26:49 +0530 Subject: [PATCH] fix(torch): densify ragged knn_search output in PointTransformer's knn_batch knn_search (open3d.ml.torch.ops) documents that it returns fewer than k neighbors per query whenever there are fewer than k candidate points -- its neighbors_index output is then ragged (not a multiple of k in total length), and its neighbors_row_splits marks each query's sub-range. knn_batch ignored this and called ans.neighbors_index.reshape(-1, k) directly, which raises "shape '[-1, 16]' is invalid for input of size N" (or silently misaligns rows on a lucky divisor) whenever any query in a batch has fewer than k neighbors -- e.g. a training config with widely spaced points, or a point cloud smaller than k. FixedRadiusSearch's ragged output hits the same shape elsewhere in this repo (kpconv.py's batch_neighbors) and is already densified with open3d.ml.torch.ops.ragged_to_dense using neighbors_row_splits; apply the same op to knn_search's output here. knn_batch's result is used by queryandgroup to index directly into points/feat, so padding cannot use ragged_to_dense's out-of-range default_value as-is (points.shape[0] would raise IndexError instead). After densifying, replace padded slots with that query's own first neighbor (a valid, already in-range index), and fall back to point 0 only for the pathological case of a query with zero neighbors at all. Distances are padded the same way, defaulting to 0.0 in that same all-empty edge case so no inf reaches downstream softmax/attention weights. Fixes #683 Co-Authored-By: Claude Sonnet 5 --- ml3d/torch/models/point_transformer.py | 39 ++++++++++++++++++++++---- tests/test_models_torch.py | 32 +++++++++++++++++++++ 2 files changed, 65 insertions(+), 6 deletions(-) diff --git a/ml3d/torch/models/point_transformer.py b/ml3d/torch/models/point_transformer.py index 130d877b..be6b50f0 100644 --- a/ml3d/torch/models/point_transformer.py +++ b/ml3d/torch/models/point_transformer.py @@ -5,7 +5,7 @@ import open3d.core as o3c from sklearn.neighbors import KDTree -from open3d.ml.torch.ops import knn_search +from open3d.ml.torch.ops import knn_search, ragged_to_dense from .base_model import BaseModel from ...utils import MODEL @@ -728,13 +728,40 @@ def knn_batch(points, points_row_splits=points_row_splits, queries_row_splits=queries_row_splits, return_distances=True) + + # knn_search returns fewer than k neighbors per query whenever there are + # fewer than k candidate points, so neighbors_index is ragged (its total + # length need not be a multiple of k) and a flat reshape(-1, k) can raise + # or silently misalign rows. Densify to (num_queries, k) using + # neighbors_row_splits, the same way FixedRadiusSearch's ragged output is + # densified elsewhere in this repo (see kpconv.py), then fill any padded + # slots by repeating that query's own first neighbor -- callers such as + # queryandgroup index directly into `points`/`feat` with this result, so + # padding must stay a valid in-range index rather than a sentinel. + n_queries = queries.shape[0] + pad_sentinel = points.shape[0] + idx = ragged_to_dense(ans.neighbors_index.reshape(-1, 1), + ans.neighbors_row_splits, k, + torch.tensor([pad_sentinel], + dtype=torch.int64)).reshape( + n_queries, k) + first_neighbor = idx[:, 0:1].expand(-1, k) + idx = torch.where(idx == pad_sentinel, first_neighbor, idx) + # A query with zero neighbors at all (pathological, but possible for a + # point cloud with fewer than k points) leaves its whole row at the + # sentinel even after the fill above; clamp those to point 0 so every + # index stays valid for points/feat. + idx = torch.where(idx == pad_sentinel, torch.zeros_like(idx), idx) if return_distances: - return ans.neighbors_index.reshape( - -1, - k).long().to(device), ans.neighbors_distance.reshape(-1, - k).to(device) + dist = ragged_to_dense( + ans.neighbors_distance.reshape(-1, 1), ans.neighbors_row_splits, + k, torch.tensor([float("inf")])).reshape(n_queries, k) + first_distance = dist[:, 0:1].expand(-1, k) + dist = torch.where(torch.isinf(dist), first_distance, dist) + dist = torch.where(torch.isinf(dist), torch.zeros_like(dist), dist) + return idx.long().to(device), dist.to(device) else: - return ans.neighbors_index.reshape(-1, k).long().to(device) + return idx.long().to(device) def interpolation(points, diff --git a/tests/test_models_torch.py b/tests/test_models_torch.py index 9c0219d7..8b58183a 100644 --- a/tests/test_models_torch.py +++ b/tests/test_models_torch.py @@ -512,6 +512,38 @@ def run(device): assert out_cpu.shape[1] == 13 +def test_pointtransformer_knn_batch_ragged_neighbors(): + """knn_search returns fewer than k neighbors per query whenever there are + fewer candidate points than k (see docstring of open3d.ml.torch.ops.knn_search). + knn_batch used to reshape(-1, k) that ragged result directly, which raises or + misaligns rows whenever any query in the batch has fewer than k neighbors. + + Regression test for https://github.com/isl-org/Open3D-ML/issues/683. + """ + from open3d.ml.torch.models.point_transformer import knn_batch + + torch.manual_seed(0) + k = 16 + n_points = 5 # fewer than k: every query gets fewer than k neighbors + points = torch.rand((n_points, 3), dtype=torch.float32) + queries = points + points_row_splits = torch.tensor([0, n_points], dtype=torch.int64) + queries_row_splits = torch.tensor([0, n_points], dtype=torch.int64) + + idx, dist = knn_batch(points, queries, k, points_row_splits, + queries_row_splits, return_distances=True) + + assert idx.shape == (n_points, k) + assert dist.shape == (n_points, k) + assert (idx >= 0).all() and (idx < n_points).all(), ( + "every neighbor index must stay valid for points/feat lookups") + + idx_no_dist = knn_batch(points, queries, k, points_row_splits, + queries_row_splits, return_distances=False) + assert idx_no_dist.shape == (n_points, k) + assert (idx_no_dist >= 0).all() and (idx_no_dist < n_points).all() + + @skip_no_pytorch_ops @skip_no_accelerator def test_pvcnn_torch():