From e6c5f1ef9709145b421463a9eb9407bb6c9672ce Mon Sep 17 00:00:00 2001 From: Austin Garrett Date: Mon, 26 Mar 2018 18:56:52 -0400 Subject: [PATCH] Try some (failed) GPU optimization. ~97% of PointCNN is spent in KNN. --- pointcnn/context.py | 3 ++- pointcnn/core.py | 39 +++++++++++++++++++++------------------ pointcnn/util.py | 44 ++++++++++++++++++++++++++++++++------------ requirements.txt | 1 + 4 files changed, 56 insertions(+), 31 deletions(-) diff --git a/pointcnn/context.py b/pointcnn/context.py index e6f42b6..c26bc1f 100644 --- a/pointcnn/context.py +++ b/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 diff --git a/pointcnn/core.py b/pointcnn/core.py index 91c7cb1..83c7028 100644 --- a/pointcnn/core.py +++ b/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) diff --git a/pointcnn/util.py b/pointcnn/util.py index 4c8d9a7..e29eb5f 100644 --- a/pointcnn/util.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index aebca71..f463940 100644 --- a/requirements.txt +++ b/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