Radius serach (#27)

* Optimise radius search

* remove debug logs

Author: nicolas-chaulet

CHANGELOG.md

@@ -0,0 +1,4 @@
+# Unreleased
+
+- ball query returns squared distance instead of distance
+- leaner Point Cloud struct that avoids copying data

cpu/include/cloud.h

@@ -31,48 +31,29 @@
 
 template <typename scalar_t> struct PointCloud
 {
-    struct PointXYZ
-    {
-        scalar_t x, y, z;
-    };
-
-    std::vector<PointXYZ> pts;
-
     void set(const std::vector<scalar_t>& new_pts)
     {
-        pts.clear();
-        pts.resize(new_pts.size() / 3);
-        for (unsigned int i = 0; i < new_pts.size(); i++)
-        {
-            if (i % 3 == 0)
-            {
-                PointXYZ point;
-                point.x = new_pts[i];
-                point.y = new_pts[i + 1];
-                point.z = new_pts[i + 2];
-                pts[i / 3] = point;
-            }
-        }
+        pts = new_pts.data();
+        length = new_pts.size() / 3;
     }
     void set_batch(const std::vector<scalar_t>& new_pts, int begin, int end)
     {
-        int size = end - begin;
-        pts.clear();
-        pts.resize(size);
-        for (int i = 0; i < size; i++)
-        {
-            PointXYZ point;
-            point.x = new_pts[3 * (begin + i)];
-            point.y = new_pts[3 * (begin + i) + 1];
-            point.z = new_pts[3 * (begin + i) + 2];
-            pts[i] = point;
-        }
+        pts = new_pts.data();
+        int start = begin * 3;
+        pts += start;
+        length = (end - begin);
     }
 
     // Must return the number of data points
     inline size_t kdtree_get_point_count() const
     {
-        return pts.size();
+        return get_point_count();
+    }
+
+    // Must return the number of data points
+    inline size_t get_point_count() const
+    {
+        return length;
     }
 
     // Returns the dim'th component of the idx'th point in the class:
@@ -82,11 +63,11 @@ template <typename scalar_t> struct PointCloud
     inline scalar_t kdtree_get_pt(const size_t idx, const size_t dim) const
     {
         if (dim == 0)
-            return pts[idx].x;
+            return pts[idx * 3];
         else if (dim == 1)
-            return pts[idx].y;
+            return pts[idx * 3 + 1];
         else
-            return pts[idx].z;
+            return pts[idx * 3 + 2];
     }
 
     // Optional bounding-box computation: return false to default to a standard
@@ -98,4 +79,29 @@ template <typename scalar_t> struct PointCloud
     {
         return false;
     }
+
+    const scalar_t* get_point_ptr(const int i) const
+    {
+        return pts + i * 3;
+    }
+
+    std::array<scalar_t, 3> operator[](const size_t index) const
+    {
+        return {pts[index * 3], pts[index * 3 + 1], pts[index * 3 + 2]};
+    }
+
+private:
+    const scalar_t* pts;
+    size_t length;
 };
+
+template <typename scalar_t>
+inline std::ostream& operator<<(std::ostream& os, const PointCloud<scalar_t>& P)
+{
+    for (size_t i = 0; i < P.get_point_count(); i++)
+    {
+        auto p = P[i];
+        os << "[" << p[0] << ", " << p[1] << ", " << p[2] << "];";
+    }
+    return os;
+}

cpu/src/neighbors.cpp

@@ -2,6 +2,7 @@
 // Taken from https://github.com/HuguesTHOMAS/KPConv
 
 #include "neighbors.h"
+#include <chrono>
 #include <random>
 
 template <typename scalar_t>
@@ -29,6 +30,7 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
     // CLoud variable
     PointCloud<scalar_t> pcd;
     pcd.set(supports);
+
     // Cloud query
     PointCloud<scalar_t> pcd_query;
     pcd_query.set(queries);
@@ -50,17 +52,17 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
     // Search params
     nanoflann::SearchParams search_params;
     search_params.sorted = sorted;
-    std::vector<std::vector<std::pair<size_t, scalar_t>>> list_matches(pcd_query.pts.size());
+    auto num_query_points = pcd_query.get_point_count();
+    std::vector<std::vector<std::pair<size_t, scalar_t>>> list_matches(num_query_points);
 
-    for (auto& p0 : pcd_query.pts)
+    for (size_t i = 0; i < num_query_points; i++)
     {
         // Find neighbors
-        scalar_t query_pt[3] = {p0.x, p0.y, p0.z};
         list_matches[i0].reserve(max_count);
         std::vector<std::pair<size_t, scalar_t>> ret_matches;
 
-        const size_t nMatches =
-            index->radiusSearch(&query_pt[0], search_radius, ret_matches, search_params);
+        const size_t nMatches = index->radiusSearch(pcd_query.get_point_ptr(i), search_radius,
+                                                    ret_matches, search_params);
         if (nMatches == 0)
             list_matches[i0] = {std::make_pair(0, -1)};
         else
@@ -164,10 +166,11 @@ int batch_nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppo
     PointCloud<scalar_t> current_cloud;
     PointCloud<scalar_t> query_pcd;
     query_pcd.set(queries);
-    vector<vector<pair<size_t, scalar_t>>> all_inds_dists(query_pcd.pts.size());
+    auto num_query_points = query_pcd.get_point_count();
+    vector<vector<pair<size_t, scalar_t>>> all_inds_dists(num_query_points);
 
     // Tree parameters
-    nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
+    nanoflann::KDTreeSingleIndexAdaptorParams tree_params(15 /* max leaf */);
 
     // KDTree type definition
     typedef nanoflann::KDTreeSingleIndexAdaptor<
@@ -178,34 +181,33 @@ int batch_nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppo
     current_cloud.set_batch(supports, s_batches[b], s_batches[b + 1]);
     std::unique_ptr<my_kd_tree_t> index(new my_kd_tree_t(3, current_cloud, tree_params));
     index->buildIndex();
+
     // Search neigbors indices
     // ***********************
     // Search params
     nanoflann::SearchParams search_params;
     search_params.sorted = sorted;
-    for (auto& p0 : query_pcd.pts)
+    std::chrono::microseconds duration_search(0);
+    for (size_t i = 0; i < num_query_points; i++)
     {
         // Check if we changed batch
-
         if (i0 == q_batches[b + 1] && b < (int)s_batches.size() - 1 &&
             b < (int)q_batches.size() - 1)
         {
             // Change the points
             b++;
-            current_cloud.pts.clear();
             if (s_batches[b] < s_batches[b + 1])
                 current_cloud.set_batch(supports, s_batches[b], s_batches[b + 1]);
 
-            // Build KDTree of the current element of the batch
             index.reset(new my_kd_tree_t(3, current_cloud, tree_params));
             index->buildIndex();
         }
 
         // Find neighboors
         std::vector<std::pair<size_t, scalar_t>> ret_matches;
         ret_matches.reserve(max_count);
-        scalar_t query_pt[3] = {p0.x, p0.y, p0.z};
-        size_t nMatches = index->radiusSearch(query_pt, r2, ret_matches, search_params);
+        size_t nMatches =
+            index->radiusSearch(query_pcd.get_point_ptr(i), r2, ret_matches, search_params);
 
         // Shuffle if needed
         if (!sorted)
@@ -225,8 +227,8 @@ int batch_nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppo
     const int token = -1;
     if (mode == 0)
     {
-        neighbors_indices.resize(query_pcd.pts.size() * max_count);
-        dists.resize(query_pcd.pts.size() * max_count);
+        neighbors_indices.resize(query_pcd.get_point_count() * max_count);
+        dists.resize(query_pcd.get_point_count() * max_count);
         i0 = 0;
         b = 0;
 
@@ -319,14 +321,15 @@ void nanoflann_knn_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppor
     // Search neigbors indices
     // ***********************
     size_t current_pos = 0;
-    for (auto& p0 : pcd_query.pts)
+    auto num_query_points = pcd_query.get_point_count();
+    for (size_t i = 0; i < num_query_points; i++)
     {
         // Find neighbors
-        scalar_t query_pt[3] = {p0.x, p0.y, p0.z};
         std::vector<size_t> ret_index(k);
         std::vector<scalar_t> out_dist_sqr(k);
 
-        const size_t nMatches = index->knnSearch(&query_pt[0], k, &ret_index[0], &out_dist_sqr[0]);
+        const size_t nMatches =
+            index->knnSearch(pcd_query.get_point_ptr(i), k, &ret_index[0], &out_dist_sqr[0]);
         for (size_t i = 0; i < nMatches; i++)
         {
             neighbors_indices[i + current_pos] = ret_index[i];

test/speed_radius.py

@@ -0,0 +1,44 @@
+import torch
+import os
+import sys
+import numpy.testing as npt
+import numpy as np
+from sklearn.neighbors import KDTree
+import unittest
+import time
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
+
+from torch_points_kernels import ball_query
+
+
+class TestRadiusSpeed(unittest.TestCase):
+    def test_speed(self):
+        start = time.time()
+        a = torch.randn(50000, 3).to(torch.float)
+        b = torch.randn(10000, 3).to(torch.float)
+        batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] + [1 for i in range(a.shape[0] // 2, a.shape[0])])
+        batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
+        R = 1
+        samples = 50
+
+        idx, dist = ball_query(R, samples, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True)
+        idx1, dist = ball_query(R, samples, a, b, mode="PARTIAL_DENSE", batch_x=batch_a, batch_y=batch_b, sort=True)
+        print(time.time() - start)
+        torch.testing.assert_allclose(idx1, idx)
+
+        self.assertEqual(idx.shape[0], b.shape[0])
+        self.assertEqual(dist.shape[0], b.shape[0])
+        self.assertLessEqual(idx.max().item(), len(batch_a))
+
+        # # Comparison to see if we have the same result
+        # tree = KDTree(a.detach().numpy())
+        # idx3_sk = tree.query_radius(b.detach().numpy(), r=R)
+        # i = np.random.randint(len(batch_b))
+        # for p in idx[i].detach().numpy():
+        #     if p >= 0 and p < len(batch_a):
+        #         assert p in idx3_sk[i]
+
+if __name__ == "__main__":
+    unittest.main()

test/test_ballquerry.py

@@ -70,7 +70,7 @@ def test_simple_gpu(self):
         dist2 = dist2.detach().cpu().numpy()
 
         idx_answer = np.asarray([[1, -1]])
-        dist2_answer = np.asarray([[0.100, -1.0000]]).astype(np.float32)
+        dist2_answer = np.asarray([[0.0100, -1.0000]]).astype(np.float32)
 
         npt.assert_array_almost_equal(idx, idx_answer)
         npt.assert_array_almost_equal(dist2, dist2_answer)
@@ -88,7 +88,7 @@ def test_simple_cpu(self):
         dist2 = dist2.detach().cpu().numpy()
 
         idx_answer = np.asarray([[1, -1]])
-        dist2_answer = np.asarray([[0.100, -1.0000]]).astype(np.float32)
+        dist2_answer = np.asarray([[0.0100, -1.0000]]).astype(np.float32)
 
         npt.assert_array_almost_equal(idx, idx_answer)
         npt.assert_array_almost_equal(dist2, dist2_answer)

torch_points_kernels/torchpoints.py

@@ -154,8 +154,6 @@ def ball_query_dense(radius, nsample, xyz, new_xyz, batch_xyz=None, batch_new_xy
         ind, dist = tpcuda.ball_query_dense(new_xyz, xyz, radius, nsample)
     else:
         ind, dist = tpcpu.dense_ball_query(new_xyz, xyz, radius, nsample, mode=0, sorted=sort)
-    positive = dist > 0
-    dist[positive] = torch.sqrt(dist[positive])
     return ind, dist
 
 
@@ -167,8 +165,6 @@ def ball_query_partial_dense(radius, nsample, x, y, batch_x, batch_y, sort=False
         ind, dist = tpcuda.ball_query_partial_dense(x, y, batch_x, batch_y, radius, nsample)
     else:
         ind, dist = tpcpu.batch_ball_query(x, y, batch_x, batch_y, radius, nsample, mode=0, sorted=sort)
-    positive = dist > 0
-    dist[positive] = torch.sqrt(dist[positive])
     return ind, dist
 
 
@@ -200,7 +196,7 @@ def ball_query(
     Returns:
         idx: (npoint, nsample) or (B, npoint, nsample) [dense] It contains the indexes of the element within x at radius distance to y
         dist: (N, nsample) or (B, npoint, nsample)  Default value: -1.
-                 It contains the distance of the element within x at radius distance to y
+                 It contains the squared distance of the element within x at radius distance to y
     """
     if mode is None:
         raise Exception('The mode should be defined within ["partial_dense | dense"]')