Try some (failed) GPU optimization. ~97% of PointCNN is spent in KNN.

pointcnn/context.py

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

pointcnn/core.py

@@ -8,10 +8,10 @@ import matplotlib.pyplot as plt
 
 try:
     from .util import knn_indices_func, MLP, BatchNorm, endchannels
-    # from .context import timed
+    from .context import timed
 except SystemError:
     from util import knn_indices_func, MLP, BatchNorm, endchannels
-    # from context import timed
+    from context import timed
 
 class XConv(nn.Module):
     """
@@ -42,13 +42,14 @@ class XConv(nn.Module):
         self.N_rep = N_rep
 
         # Additional processing layers
-        self.pts_batchnorm = nn.BatchNorm2d(N_rep, momentum = 0.9)
+        self.pts_batchnorm = BatchNorm(2, D, momentum = 0.9)
+        # self.pts_batchnorm = BatchNorm(BatchNorm())
 
         # Main dense linear layers
         self.mlp_lift = MLP(np.around(np.geomspace(D, self.C_lifted, num = mlp_width)).astype(int))
-        self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1))
+        self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1).cuda())
         self.mlp = MLP(np.around(np.geomspace(N_neighbors, N_neighbors)).astype(int))  # Somehow, original code has K x K.
-        self.end_conv = endchannels(nn.Conv2d(C_lifted + C_in, C_out, (N_neighbors, 1), groups = C_out))
+        self.end_conv = endchannels(nn.Conv2d(C_lifted + C_in, C_out, (N_neighbors, 1), groups = C_out).cuda())
 
         # Params for kernel initialization.
         self.K = nn.Parameter(torch.FloatTensor(C_out, C_in + self.C_lifted, N_neighbors))
@@ -136,6 +137,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
@@ -150,14 +152,14 @@ 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.cpu(), P.cpu(), N_neighbors)
-        P_regional = self.select_region(P, P_idx)
-        if False:
+        P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors).cuda()  # This step takes ~97% of the time.
+        P_regional = self.select_region(P, P_idx)                           # Prime target for optimization: KNN on GPU.
+        if True:
             # Draw neighborhood points, for debugging.
             t = 23
             n = 3
-            test_point = ps[n,t,:].data.numpy()
-            neighborhood = P_regional[n,t,:,:].data.numpy()
+            test_point = ps[n,t,:].cpu().data.numpy()
+            neighborhood = P_regional[n,t,:,:].cpu().data.numpy()
             plt.scatter(P[n][:,0], P[n][:,1])
             plt.scatter(test_point[0], test_point[1], s = 100, c = 'green')
             plt.scatter(neighborhood[:,0], neighborhood[:,1], s = 100, c = 'red')
@@ -190,19 +192,20 @@ if __name__ == "__main__":
         num_points = 1000
         N_rep = 50
         D = 2
-        C_in = 64
-        C_out = 128
-        N_neighbors = 10
+        C_in = 128
+        C_out = 256
+        N_neighbors = 5
 
-        model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, knn_indices_func)
+        model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, 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()
 
-        out = model(test_ps, test_P, test_F)
+        for _ in range(10):
+            out = model(test_ps, test_P, test_F)

pointcnn/util.py

@@ -1,9 +1,24 @@
+import time
+
 import torch
 import torch.nn as nn
 
 import numpy as np
 from sklearn.neighbors import NearestNeighbors
 
+try:
+    from .context import pytorch_knn_cuda
+except SystemError:
+    from context import pytorch_knn_cuda
+
+def endchannels(f, make_contiguous = False):
+    def wrapped_func(x):
+        if make_contiguous:
+            return torch.transpose(f(torch.transpose(x, 1, -1).contiguous()), -1, 1)
+        else:
+            return torch.transpose(f(torch.transpose(x, 1, -1)), -1, 1)
+    return wrapped_func
+
 class BatchNorm(nn.Module):
     """
     PyTorch Linear layers transform shape in the form (N,*,in_features) ->
@@ -26,14 +41,6 @@ class BatchNorm(nn.Module):
 
         self.forward = endchannels(self.bn, make_contiguous = True)
 
-def endchannels(f, make_contiguous = False):
-    def wrapped_func(x):
-        if make_contiguous:
-            return torch.transpose(f(torch.transpose(x, 1, -1).contiguous()), -1, 1)
-        else:
-            return torch.transpose(f(torch.transpose(x, 1, -1)), -1, 1)
-    return wrapped_func
-
 def MLP(layer_sizes, activation_layer = nn.ReLU(), batch_norm = True):
     """
     Creates a fully connected MLP of arbitrary depth.
@@ -95,7 +102,6 @@ def knn_indices_func(ps, P, k):
     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:]
@@ -105,6 +111,20 @@ def knn_indices_func(ps, P, k):
     ], axis = 0)
     return torch.from_numpy(region_idx)
 
+def knn_indices_func_gpu(ps, P, k):
+
+    def single_batch_knn(p, P_particular):
+        nbrs_f = pytorch_knn_cuda.KNearestNeighbor(k + 1)
+        # knn_cuda(k + 1, )
+        indices = nbrs_f(P_particular, p)[0]
+        return indices[:,1:]
+
+    region_idx = torch.stack([
+        single_batch_knn(p, P[n]) for n, p in enumerate(ps)
+    ], dim = 0)
+    return region_idx
+
+
 if __name__ == "__main__":
     from torch.autograd import Variable
     N_rep = 100
@@ -114,7 +134,7 @@ if __name__ == "__main__":
     test_P  = np.random.rand(N,num_points,D).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_ps = Variable(torch.from_numpy(test_ps))
-    out = knn_indices_func(test_ps, test_P, 5)
+    test_P = Variable(torch.from_numpy(test_P)).cuda()
+    test_ps = Variable(torch.from_numpy(test_ps)).cuda()
+    out = knn_indices_func_gpu(test_ps, test_P, 5)
     print(out)

requirements.txt

@@ -1,5 +1,6 @@
 nose
 sphinx
 numpy
+libKMCUDA
 http://download.pytorch.org/whl/cu80/torch-0.3.1-cp35-cp35m-linux_x86_64.whl
 torchvision