From a0f3d762feb28b4a1495335cd1c362a470fb2d6a Mon Sep 17 00:00:00 2001 From: Austin Garrett Date: Fri, 23 Mar 2018 17:22:49 -0400 Subject: [PATCH] Finish about half of vectorizing XConv. --- pointcnn/core.py | 160 ++++++++++++++++++++++++++++++++++------------- 1 file changed, 115 insertions(+), 45 deletions(-) diff --git a/pointcnn/core.py b/pointcnn/core.py index 5c39123..692a8b6 100644 --- a/pointcnn/core.py +++ b/pointcnn/core.py @@ -12,24 +12,63 @@ 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 MLP(layer_sizes, activation_func = nn.ReLU()): + 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_func: Activation function to be applied in between 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() - return nn.Sequential(*[ - nn.Sequential(nn.Linear(C_in, C_out), activation_func) - for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:]) - ]) + 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:]) + ]) class XConv(nn.Module): + """ + Vectorized pointwise convolution. + """ - def __init__(self, C_in, C_out, D, N_neighbors, C_lifted = None, mlp_width = 4): + def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None, mlp_width = 4): """ :param C_in: Input dimension of the points' features. :param C_out: Output dimension of the representative point features. @@ -50,9 +89,19 @@ class XConv(nn.Module): self.D = D self.N_neighbors = 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.N_rep = N_rep + # Additional processing layers + self.pts_batchnorm = nn.BatchNorm2d(N_rep, momentum = 0.9) + + # 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. + ) + + # Params for kernel initialization. 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) @@ -60,38 +109,42 @@ class XConv(nn.Module): def forward(self, p, P, F): """ Applies XConv to the input data. - :type p: FloatTensor (N, D) - :type P: FloatTensor (N, N_neighbors, D) - :type F: FloatTensor (N, N_neighbors, C_in) + :type p: FloatTensor (N, N_rep, D) + :type P: FloatTensor (N, N_rep, N_neighbors, D) + :type F: FloatTensor (N, N_rep, N_neighbors, C_in) :rtype: FloatTensor (TODO: shape) :param p: Representative point :param P: Regional point cloud such that F[:,p_idx,:] is the feature associated with P[:,p_idx,:] :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. + assert(p.size()[0] == P.size()[0] == F.size()[0]) # Check N is equal. + assert(p.size()[1] == P.size()[1] == F.size()[1] == self.N_rep) # Check N_rep is equal. + assert(P.size()[2] == F.size()[2] == self.N_neighbors) # Check N_neighbors is equal. + assert(p.size()[2] == P.size()[3] == self.D) # Check D is equal. + assert(F.size()[3] == 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. - F_lifted = self.mlp_lift(P_loc) # Individually lift each point into C_lifted dim space. - F_cat = torch.cat((F_lifted, F), 2) # Cat F_lifted and F, to size (N, K, C_lifted + C_in). - X = self.mlp(P_loc) # Learn the (N, K, K) X-transformation matrix. - 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 = torch.squeeze(nn.functional.conv1d( # Finally, typical convolution between K and F_X. + 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). + 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 - )) + ) - return F_p + return torch.squeeze(F_p, dim = 2) class PointCNN(nn.Module): - def __init__(self, C_in, C_out, D, N_neighbors, r_indices_func, C_lifted = None, mlp_width = 4): + def __init__(self, C_in, C_out, D, N_neighbors, N_rep, r_indices_func, C_lifted = None, mlp_width = 4): """ :param C_in: Input dimension of the points' features. :param C_out: Output dimension of the representative point features. @@ -120,7 +173,7 @@ class PointCNN(nn.Module): C_lifted = C_in # Not optimal? self.r_indices_func = r_indices_func - self.x_conv = XConv(C_in, C_out, D, N_neighbors, C_lifted, mlp_width) + self.x_conv = XConv(C_in, C_out, D, N_neighbors, N_rep, C_lifted, mlp_width) def select_region(self, P, P_idx): """ @@ -145,7 +198,7 @@ class PointCNN(nn.Module): :type ps: FloatTensor (N, *, D) :type P: FloatTensor (N, N_neighbors, D) :type F: FloatTensor (N, N_neighbors, C_in) - :rtype: FloatTensor (TODO: shape) + :rtype: FloatTensor (N, N_rep, D) :param ps: Representative points :param P: Regional point cloud such that F[:,p_idx,:] is the feature associated with P[:,p_idx,:] :param F: Regional features such that P[:,p_idx,:] is the feature associated with F[:,p_idx,:] @@ -161,23 +214,40 @@ class PointCNN(nn.Module): if __name__ == "__main__": np.random.seed(0) - N = 4 - num_points = 15000 - N_rep = 7500 - D = 3 - C_in = 8 - C_out = 32 - N_neighbors = 5 + TESTING = XConv - model = PointCNN(C_in, C_out, D, N_neighbors, knn_indices_func).cuda() + if TESTING == XConv: + N = 4 + N_rep = 100 + D = 3 + C_in = 8 + C_out = 32 + N_neighbors = 10 - 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,:] + 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() + out = model(p, P, F) - 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() + elif TESTING == PointCNN: + N = 4 + num_points = 15000 + N_rep = 7500 + D = 3 + C_in = 8 + C_out = 32 + N_neighbors = 5 - out = model(test_ps, test_P, test_F) + 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)).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)