Add a few tests. Debug KNN.

pointcnn/context.py

@@ -1,4 +1,4 @@
 import sys, os
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
-from lib import pytorch_knn_cuda
+from lib import timed

pointcnn/core.py

@@ -1,11 +1,17 @@
+import time
+
 import torch
 import torch.nn as nn
 from torch.autograd import Variable
 import numpy as np
-from sklearn.neighbors import NearestNeighbors
 
-from context import pytorch_knn_cuda
-KNN = pytorch_knn_cuda.KNearestNeighbor
+try:
+    from .util import knn_indices_func
+    from .context import timed
+except SystemError:
+    from util import knn_indices_func
+    from context import timed
+
 
 def MLP(layer_sizes, activation_func = nn.ReLU()):
     """
@@ -32,17 +38,22 @@ class XConv(nn.Module):
         :param C_lifted: Dimensionality of lifted point features.
         :param mlp_width: Number of hidden layers in MLPs.
         """
+        super(XConv, self).__init__()
 
         if C_lifted == None:
             C_lifted = C_in  # Not optimal?
 
-        super(XConv, self).__init__()
-        self.N_neighbors = N_neighbors
-        self.D = D
-        self.mlp_lift = MLP(np.around(np.geomspace(D, C_lifted, num = mlp_width)).astype(int))
-        self.mlp = MLP(np.floor(np.geomspace(D, N_neighbors)).astype(int))
+        if __debug__:
+            # Only needed for assertions.
+            self.C_in = C_in
+            self.C_lifted = C_lifted
+            self.D = D
+            self.N_neighbors = N_neighbors
 
-        self.K = nn.Parameter(torch.FloatTensor(C_out, C_in + C_lifted, N_neighbors))
+        self.mlp_lift = MLP(np.around(np.geomspace(D, self.C_lifted, num = mlp_width)).astype(int))
+        self.mlp = MLP(np.around(np.geomspace(D, N_neighbors)).astype(int))
+
+        self.K = nn.Parameter(torch.FloatTensor(C_out, C_in + self.C_lifted, N_neighbors))
         stdv = 1. / np.sqrt(N_neighbors)
         self.K.data.uniform_(-stdv, stdv)
 
@@ -58,10 +69,10 @@ class XConv(nn.Module):
         :param F: Regional features such that P[:,p_idx,:] is the feature associated with F[:,p_idx,:]
         :return: Features aggregated into point p.
         """
-
         assert(p.size()[0] == P.size()[0] == F.size()[0])       # Check N is equal.
         assert(P.size()[1] == F.size()[1] == self.N_neighbors)  # Check N_neighbors is equal.
         assert(p.size()[1] == P.size()[2] == self.D)            # Check D is equal.
+        assert(F.size()[2] == self.C_in)                        # Check C_in is equal.
 
         N = len(P)
         P_loc = P - torch.unsqueeze(p, 1)               # Move P to local coordinate system of p.
@@ -71,11 +82,12 @@ class XConv(nn.Module):
         F_X = torch.stack([                             # Weight and permute F_cat with the learned X.
             torch.mm(X[n], F_cat[n]) for n in range(N)
         ], dim = 0)
-        F_p = nn.functional.conv1d(                     # Finally, typical convolution between K and F_X.
+        F_p = torch.squeeze(nn.functional.conv1d(       # Finally, typical convolution between K and F_X.
             torch.transpose(F_X, 1, 2),
             self.K
-        )
-        return F_p.view(N, -1)
+        ))
+
+        return F_p
 
 class PointCNN(nn.Module):
 
@@ -103,6 +115,7 @@ class PointCNN(nn.Module):
         :param mlp_width: Number of hidden layers in MLPs.
         """
         super(PointCNN, self).__init__()
+
         if C_lifted == None:
             C_lifted = C_in  # Not optimal?
 
@@ -123,6 +136,7 @@ class PointCNN(nn.Module):
         ], dim = 0)
         return regions
 
+    @timed.timed
     def forward(self, ps, P, F):
         """
         Given a set of representative points, a point cloud, and its
@@ -137,59 +151,33 @@ class PointCNN(nn.Module):
         :param F: Regional features such that P[:,p_idx,:] is the feature associated with F[:,p_idx,:]
         :return:
         """
-        P_idx = self.r_indices_func(ps, P, N_neighbors)
+        P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors).cuda()
         inp_regions = torch.stack([
             self.x_conv(p, self.select_region(P, P_idx[:,n]), self.select_region(F, P_idx[:,n]))
             for n, p in enumerate(torch.unbind(ps, dim = 1))
         ], dim = 1)
         return inp_regions
 
-def knn_indices_func(ps, P, k):
-    """
-    Indexing function based on K-Nearest Neighbors search.
-    :type p: FloatTensor (N, D)
-    :type P: FloatTensor (N, *, D)
-    :rtype: FloatTensor (N, N_neighbors)
-    :param p: Representative point
-    :param P: Point cloud to get indices from
-    :return: Array of indices, P_idx, into P such that P[P_idx] is the set
-    of points in the "region" around p.
-    """
-    ps = ps.data.numpy()
-    P = P.data.numpy()
-
-    def single_batch_knn(p, P_particular):
-        nbrs = NearestNeighbors(k, algorithm = "ball_tree").fit(p)
-        indices = nbrs.kneighbors(P_particular)[1]
-        return indices
-
-    region_idx = np.stack([
-        single_batch_knn(p, P[n]) for n, p in enumerate(ps)
-    ], axis = 0)
-    return torch.from_numpy(region_idx)
-
 if __name__ == "__main__":
     np.random.seed(0)
 
     N = 4
-    num_points = 5000
-    N_rep = 1000
+    num_points = 15000
+    N_rep = 7500
     D = 3
     C_in = 8
     C_out = 32
-    N_neighbors = 100
+    N_neighbors = 5
 
-    model = PointCNN(C_in, C_out, D, N_neighbors, knn_indices_func)
+    model = PointCNN(C_in, C_out, D, N_neighbors, knn_indices_func).cuda()
 
     test_P  = np.random.rand(N,num_points,D).astype(np.float32)
     test_F  = np.random.rand(N,num_points,C_in).astype(np.float32)
     idx = np.random.choice(test_P.shape[1], N_rep, replace = False)
     test_ps = test_P[:,idx,:]
 
-    test_P = Variable(torch.from_numpy(test_P))
-    test_F = Variable(torch.from_numpy(test_F))
-    test_ps = Variable(torch.from_numpy(test_ps))
+    test_P = Variable(torch.from_numpy(test_P)).cuda()
+    test_F = Variable(torch.from_numpy(test_F)).cuda()
+    test_ps = Variable(torch.from_numpy(test_ps)).cuda()
 
-    print(test_P)
     out = model(test_ps, test_P, test_F)
-    print(out)

pointcnn/util.py

@@ -1,4 +1,6 @@
 import torch
+import numpy as np
+from sklearn.neighbors import NearestNeighbors
 
 def apply_along_dim(xs, f, dim):
     """
@@ -21,6 +23,33 @@ def zipwith_matmul(xs, ys):
     N = len(xs)
     return torch.stack([torch.mm(xs[i], ys[i]) for i in range(N)], dim = 0)
 
-foo = torch.FloatTensor(10, 5, 6)
-bar = torch.FloatTensor(10, 6, 7)
-baz = zipwith_matmul(foo, bar)
+def knn_indices_func(ps, P, k):
+    """
+    Indexing function based on K-Nearest Neighbors search.
+    :type ps: FloatTensor (N, N_rep, D)
+    :type P: FloatTensor (N, *, D)
+    :type k: int
+    :rtype: FloatTensor (N, N_rep, N_neighbors)
+    :param ps: Representative point
+    :param P: Point cloud to get indices from
+    :param k: Number of nearest neighbors to collect.
+    :return: Array of indices, P_idx, into P such that P[n][P_idx[n],:]
+    is the set k-nearest neighbors for the representative points in P[n].
+    """
+    ps = ps.data.numpy()
+    P = P.data.numpy()
+
+    def single_batch_knn(p, P_particular):
+        # p, P_particular = P_particular, p
+        nbrs = NearestNeighbors(k + 1, algorithm = "ball_tree").fit(P_particular)
+        indices = nbrs.kneighbors(p)[1]
+        return indices[:,1:]
+
+    region_idx = np.stack([
+        single_batch_knn(p, P[n]) for n, p in enumerate(ps)
+    ], axis = 0)
+    return torch.from_numpy(region_idx)
+
+if __name__ == "__main__":
+    from torch.autograd import Variable
+

tests/test_basic.py

@@ -4,7 +4,7 @@ import torch
 from torch.autograd import Variable
 import numpy as np
 
-from pointcnn.core import XConv
+from pointcnn.core import XConv, knn_indices_func
 
 class BasicTests(unittest.TestCase):
     """ Basic test cases """
@@ -19,12 +19,29 @@ class BasicTests(unittest.TestCase):
         C_out = 32
         N_neighbors = 100
 
-        model = XConv(C_in, C_out, D, N_neighbors)
-        test_p = Variable(torch.from_numpy(np.random.rand(N,D).astype(np.float32)))
-        test_P = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,D).astype(np.float32)))
-        test_F = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,C_in).astype(np.float32)))
-        test_out = model(test_p, test_P, test_F)
-        self.assertEqual(test_out.size(), (N, C_out))
+        model = XConv(C_in, C_out, D, N_neighbors).cuda()
+        p = Variable(torch.from_numpy(np.random.rand(N,D).astype(np.float32))).cuda()
+        P = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,D).astype(np.float32))).cuda()
+        F = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,C_in).astype(np.float32))).cuda()
+        out = model(p, P, F)
+        self.assertEqual(out.size(), (N, C_out))
+
+    def test_knn(self):
+        P = np.array([[[0,0],
+                       [0,0.95],
+                       [1,0],
+                       [1,1]]])
+        ps = P[:,[0,3],:]
+
+        P = Variable(torch.from_numpy(P))
+        ps = Variable(torch.from_numpy(ps))
+
+        out = knn_indices_func(ps, P, 2).numpy()
+        target = np.array([[[1,2],
+                            [2,1]]])
+
+        self.assertTrue(np.array_equal(target, out))
+
 
 if __name__ == "__main__":
     unittest.main()