From 9be3f8528a784be6503cd37ded4b03a86f8c0c6a Mon Sep 17 00:00:00 2001 From: "Austin J. Garrett" Date: Sun, 25 Mar 2018 18:32:18 -0400 Subject: [PATCH] Vectorize XConv. Tests still required. --- pointcnn/context.py | 2 +- pointcnn/core.py | 155 ++++++++++++++++---------------------------- pointcnn/util.py | 67 ++++++++++++++++++- 3 files changed, 122 insertions(+), 102 deletions(-) diff --git a/pointcnn/context.py b/pointcnn/context.py index 805dc19..e6f42b6 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 timed +# from lib import timed diff --git a/pointcnn/core.py b/pointcnn/core.py index 692a8b6..91c7cb1 100644 --- a/pointcnn/core.py +++ b/pointcnn/core.py @@ -4,71 +4,21 @@ import torch import torch.nn as nn from torch.autograd import Variable import numpy as np +import matplotlib.pyplot as plt try: - from .util import knn_indices_func - from .context import timed + from .util import knn_indices_func, MLP, BatchNorm, endchannels + # from .context import timed except SystemError: - from util import knn_indices_func - from context import timed - -class BatchNorm(nn.Module): - """ - PyTorch Linear layers transform shape in the form (N,*,in_features) -> - (N,*,out_features). BatchNorm normalizes over axis 1. Thus, BatchNorm - following a linear layer ONLY has the desired behavior if there are no - additional (*) dimensions. To get the desired behavior, we first transpose - the appropriate axis into the channel dim, then tranpsose out. - """ - - def __init__(self, D, num_features, dim = 1, *args, **kwargs): - super(BatchNorm, self).__init__() - if D == 1: - self.bn = nn.BatchNorm1d(num_features, *args, **kwargs) - elif D == 2: - self.bn = nn.BatchNorm2d(num_features, *args, **kwargs) - elif D == 3: - self.bn = nn.BatchNorm3d(num_features, *args, **kwargs) - else: - raise ValueError("Dimensionality %i not supported" % D) - - self.dim = dim - - def forward(self, x): - x = torch.transpose(x, 1, self.dim).contiguous() # Must be made contiguous for cudNN. - self.bn(x) - x = torch.transpose(x, self.dim, 1) - return x - -def MLP(layer_sizes, activation_layer = nn.ReLU(), batch_norm = True): - """ - Creates a fully connected MLP of arbitrary depth. - :param layer_sizes: Sizes of MLP hidden layers. - :param activation_layer: Activation function to be applied in between layers. - :return: Multilayer perceptron module - """ - if isinstance(layer_sizes, np.ndarray): - layer_sizes = layer_sizes.tolist() - if batch_norm: - return nn.Sequential(*[ - nn.Sequential(nn.Linear(C_in, C_out), - activation_layer, - BatchNorm(D = 2, num_features = C_out, dim = -1, momentum = 0.9) - ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:]) - ]) - else: - return nn.Sequential(*[ - nn.Sequential(nn.Linear(C_in, C_out), - activation_layer, - ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:]) - ]) + from util import knn_indices_func, MLP, BatchNorm, endchannels + # from context import timed class XConv(nn.Module): """ Vectorized pointwise convolution. """ - def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None, mlp_width = 4): + def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None, mlp_width = 2): """ :param C_in: Input dimension of the points' features. :param C_out: Output dimension of the representative point features. @@ -96,10 +46,9 @@ class XConv(nn.Module): # Main dense linear layers self.mlp_lift = MLP(np.around(np.geomspace(D, self.C_lifted, num = mlp_width)).astype(int)) - self.mlp = nn.Sequential( - # torch.Conv2d(TODO), - MLP(np.around(np.geomspace(3, N_neighbors)).astype(int)) # Somehow, original code has K x K. - ) + self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1)) + 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)) # Params for kernel initialization. self.K = nn.Parameter(torch.FloatTensor(C_out, C_in + self.C_lifted, N_neighbors)) @@ -126,23 +75,20 @@ class XConv(nn.Module): N = len(P) p_center = torch.unsqueeze(p, dim = 2) - P_local = self.pts_batchnorm(P - p_center) # Move P to local coordinate system of p. - F_lifted = self.mlp_lift(P_local) # Individually lift each point into C_lifted dim space. - F_cat = torch.cat((F_lifted, F), -1) # Cat F_lifted and F, to size (N, N_rep, N_neighbors, C_lifted + C_in). + P_local = self.pts_batchnorm(P - p_center) # Move P to local coordinate system of p. + F_lifted = self.mlp_lift(P_local) # Individually lift each point into C_lifted dim space. + F_cat = torch.cat((F_lifted, F), -1) # Cat F_lifted and F, to size (N, N_rep, N_neighbors, C_lifted + C_in). X_shape = (N, self.N_rep, N_neighbors, N_neighbors) - X = self.mlp(P_local).contiguous().view(*X_shape) # Learn the (N, K, K) X-transformation matrix. - F_X = torch.matmul(X, F_cat) # Weight and permute F_cat with the learned X. - - # CODE PAST THIS POINT BROKEN. - # TODO: Implement separable_conv2d - F_p = nn.functional.conv1d( # Finally, typical convolution between K and F_X. - torch.transpose(F_X, 1, 2), - self.K - ) - + X = self.mlp(self.mid_conv(P_local)) # Learn the (N, K, K) X-transformation matrix. + X = X.contiguous().view(*X_shape) + F_X = torch.matmul(X, F_cat) # Weight and permute F_cat with the learned X. + F_p = self.end_conv(F_X) return torch.squeeze(F_p, dim = 2) class PointCNN(nn.Module): + """ + TODO: Insert documentation + """ def __init__(self, C_in, C_out, D, N_neighbors, N_rep, r_indices_func, C_lifted = None, mlp_width = 4): """ @@ -151,7 +97,7 @@ class PointCNN(nn.Module): :param D: Spatial dimensionality of points. :param N_neighbors: Number of neighbors to convolve over. :param r_indices_func: Selector function of the type, - INP + INP ====== ps : (N, N_rep, D) Representative points P : (N, *, D) Point cloud @@ -163,7 +109,7 @@ class PointCNN(nn.Module): P[P_idx] is the set of points in the "region" around p. a representative point p and a point cloud P. From these it returns an - array of N_neighbors + array of N_neighbors :param C_lifted: Dimensionality of lifted point features. :param mlp_width: Number of hidden layers in MLPs. """ @@ -177,19 +123,19 @@ class PointCNN(nn.Module): def select_region(self, P, P_idx): """ - Selects + Selects :type P: FloatTensor (N, *, D) - :type P_idx: FloatTensor (N, N_neighbors) + :type P_idx: FloatTensor (N, N_rep, N_neighbors) :rtype P_region: FloatTensor (N_rep, N_neighbors, D) :param P: Point cloud to select regional points from :param P_idx: Indices of points in region to be selected + :return: """ regions = torch.stack([ P[n][idx,:] for n, idx in enumerate(torch.unbind(P_idx, dim = 0)) ], dim = 0) return regions - @timed.timed def forward(self, ps, P, F): """ Given a set of representative points, a point cloud, and its @@ -204,50 +150,59 @@ 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).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 + P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors) + P_regional = self.select_region(P, P_idx) + if False: + # Draw neighborhood points, for debugging. + t = 23 + n = 3 + test_point = ps[n,t,:].data.numpy() + neighborhood = P_regional[n,t,:,:].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') + plt.show() + F_regional = self.select_region(F, P_idx) + return self.x_conv(ps, P_regional, F_regional) if __name__ == "__main__": np.random.seed(0) - TESTING = XConv + TESTING = PointCNN if TESTING == XConv: N = 4 - N_rep = 100 + N_rep = 150 D = 3 C_in = 8 C_out = 32 N_neighbors = 10 - model = XConv(C_in, C_out, D, N_neighbors, N_rep).cuda() - p = Variable(torch.from_numpy(np.random.rand(N,N_rep,D).astype(np.float32))).cuda() - P = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,D).astype(np.float32))).cuda() - F = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,C_in).astype(np.float32))).cuda() + model = XConv(C_in, C_out, D, N_neighbors, N_rep) + p = Variable(torch.from_numpy(np.random.rand(N,N_rep,D).astype(np.float32))) + P = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,D).astype(np.float32))) + F = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,C_in).astype(np.float32))) + out = model(p, P, F) elif TESTING == PointCNN: N = 4 - num_points = 15000 - N_rep = 7500 - D = 3 - C_in = 8 - C_out = 32 - N_neighbors = 5 + num_points = 1000 + N_rep = 50 + D = 2 + C_in = 64 + C_out = 128 + N_neighbors = 10 - model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, knn_indices_func).cuda() + model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, knn_indices_func) 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)).cuda() - test_F = Variable(torch.from_numpy(test_F)).cuda() - test_ps = Variable(torch.from_numpy(test_ps)).cuda() + test_P = Variable(torch.from_numpy(test_P)) + test_F = Variable(torch.from_numpy(test_F)) + test_ps = Variable(torch.from_numpy(test_ps)) out = model(test_ps, test_P, test_F) diff --git a/pointcnn/util.py b/pointcnn/util.py index 0f6907f..4c8d9a7 100644 --- a/pointcnn/util.py +++ b/pointcnn/util.py @@ -1,7 +1,62 @@ import torch +import torch.nn as nn + import numpy as np from sklearn.neighbors import NearestNeighbors +class BatchNorm(nn.Module): + """ + PyTorch Linear layers transform shape in the form (N,*,in_features) -> + (N,*,out_features). BatchNorm normalizes over axis 1. Thus, BatchNorm + following a linear layer ONLY has the desired behavior if there are no + additional (*) dimensions. To get the desired behavior, we first transpose + the channel dim into the last dim, then tranpsose out. + """ + + def __init__(self, D, num_features, *args, **kwargs): + super(BatchNorm, self).__init__() + if D == 1: + self.bn = nn.BatchNorm1d(num_features, *args, **kwargs) + elif D == 2: + self.bn = nn.BatchNorm2d(num_features, *args, **kwargs) + elif D == 3: + self.bn = nn.BatchNorm3d(num_features, *args, **kwargs) + else: + raise ValueError("Dimensionality %i not supported" % D) + + 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. + :param layer_sizes: Sizes of MLP hidden layers. + :param activation_layer: Activation function to be applied in between layers. + :return: Multilayer perceptron module + """ + if isinstance(layer_sizes, np.ndarray): + layer_sizes = layer_sizes.tolist() + if batch_norm: + return nn.Sequential(*[ + nn.Sequential(nn.Linear(C_in, C_out), + activation_layer, + BatchNorm(D = 2, num_features = C_out, momentum = 0.9) + ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:]) + ]) + else: + return nn.Sequential(*[ + nn.Sequential(nn.Linear(C_in, C_out), + activation_layer, + ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:]) + ]) + def apply_along_dim(xs, f, dim): """ PyTorch analog to np.apply_along_axis. @@ -52,4 +107,14 @@ def knn_indices_func(ps, P, k): if __name__ == "__main__": from torch.autograd import Variable - + N_rep = 100 + N = 2 + num_points = 1000 + D = 3 + 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) + print(out)