diff --git a/pointcnn/context.py b/pointcnn/context.py index 7c5da79..805dc19 100644 --- a/pointcnn/context.py +++ b/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 diff --git a/pointcnn/core.py b/pointcnn/core.py index 8ffa968..5c39123 100644 --- a/pointcnn/core.py +++ b/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) diff --git a/pointcnn/util.py b/pointcnn/util.py index 2cafe98..0f6907f 100644 --- a/pointcnn/util.py +++ b/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 + diff --git a/tests/test_basic.py b/tests/test_basic.py index 8793800..2970dae 100644 --- a/tests/test_basic.py +++ b/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()