From 50b77758df53607734c67ceb490dcec9abaa7678 Mon Sep 17 00:00:00 2001 From: "Austin J. Garrett" Date: Sat, 17 Mar 2018 20:27:49 -0400 Subject: [PATCH] Implement XConv operator. Begin implementing PointCNN. --- Makefile | 2 +- pointcnn/__init__.py | 0 pointcnn/core.py | 139 +++++++++++++++++++++++++++++++++++++++++ pointcnn/pointcnn.py | 1 - pointcnn/util.py | 27 +++++++- tests/__init__.py | 0 tests/__main__.py | 6 -- tests/context.py | 2 - tests/test_advanced.py | 2 - tests/test_basic.py | 22 ++++++- 10 files changed, 186 insertions(+), 15 deletions(-) create mode 100644 pointcnn/__init__.py create mode 100644 pointcnn/core.py delete mode 100644 pointcnn/pointcnn.py create mode 100644 tests/__init__.py delete mode 100644 tests/__main__.py diff --git a/Makefile b/Makefile index d5bb6c1..3cc35d6 100644 --- a/Makefile +++ b/Makefile @@ -2,4 +2,4 @@ init: pip3 install -r requirements.txt test: - nosetests tests + python3.5 -m unittest discover diff --git a/pointcnn/__init__.py b/pointcnn/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pointcnn/core.py b/pointcnn/core.py new file mode 100644 index 0000000..a9ad4d8 --- /dev/null +++ b/pointcnn/core.py @@ -0,0 +1,139 @@ +import torch +import torch.nn as nn +from torch.autograd import Variable +import numpy as np + +def MLP(layer_sizes, activation_func = nn.ReLU()): + """ + 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. + :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:]) + ]) + +class XConv(nn.Module): + + def __init__(self, C_in, C_out, D, N_neighbors, 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. + :param D: Spatial dimensionality of points. + :param N_neighbors: Number of neighbors to convolve over. + :param C_lifted: Dimensionality of lifted point features. + :param mlp_width: Number of hidden layers in MLPs. + """ + + 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)) + + self.K = nn.Parameter(torch.FloatTensor(C_out, C_in + C_lifted, N_neighbors)) + stdv = 1. / np.sqrt(N_neighbors) + self.K.data.uniform_(-stdv, stdv) + + 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) + :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. + + 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 = 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) + +class PointCNN(nn.Module): + + def __init__(self, C_in, C_out, D, N_neighbors, 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. + :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 + ====== + p : (N, D) Representative point + P : (N, *, D) Point cloud + + OUT + ====== + P_idx : (N, N_neighbors) Array of indices into P such that + P[P_idx] is the set of points + + a representative point p and a point cloud P. From these it returns an + array of N_neighbors + :param C_lifted: Dimensionality of lifted point features. + :param mlp_width: Number of hidden layers in MLPs. + """ + super(PointCNN, self).__init__() + if C_lifted == None: + C_lifted = C_in # Not optimal? + + self.r_filter = r_filter + self.x_conv = XConv(C_in, C_out, D, N_neighbors, C_lifted, mlp_width) + + def select_region(self, P_idx, P): + """ + Selects + :type P_idx: FloatTensor (N, N_neighbors) + :type P: FloatTensor (N, *, *) + :param P_idx: Indices of points in region to be selected. + :param P: Point cloud to select regional points from. + """ + return torch.stack([ + P[n,:,:].index_select(0, idx) for n, idx in torch.unbind(P_idx, dim = 0) + ], dim = 0) + + def forward(self, ps, P, F): + """ + Given a set of representative points, a point cloud, and its + corresponding features, return a new set of representative points with + features projected from the point cloud. + :type p: FloatTensor (N, *, D) + :type P: FloatTensor (N, N_neighbors, D) + :type F: FloatTensor (N, 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: + """ + # (N, *, N_neighbors, D) + P_idx = self.r_indices_func(p, P) + inp_regions = torch.stack([ + self.x_conv(p, self.select_region(P, P_idx), self.select_region(F, P_idx)) + for p in torch.unbind(ps, dim = 1) + ], dim = 1) + return inp_regions diff --git a/pointcnn/pointcnn.py b/pointcnn/pointcnn.py deleted file mode 100644 index 2f4eb9a..0000000 --- a/pointcnn/pointcnn.py +++ /dev/null @@ -1 +0,0 @@ -# This is a file for the main PyTorch module. diff --git a/pointcnn/util.py b/pointcnn/util.py index 3177f3d..2cafe98 100644 --- a/pointcnn/util.py +++ b/pointcnn/util.py @@ -1 +1,26 @@ -# This is a file for utility functions. +import torch + +def apply_along_dim(xs, f, dim): + """ + PyTorch analog to np.apply_along_axis. + :param xs: + :param dim: + """ + return torch.stack([f(x) for x in torch.unbind(xs, dim)], dim) + +def zipwith_matmul(xs, ys): + """ + Given two lists of 2D matrices of appropriate size, zips them + together with matrix multiplication. + :param xs: + :param ys: + """ + # xs of shape [N, n, m] + # ys of shape [N, m, p] + # return shape [N, n, p] + 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) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/__main__.py b/tests/__main__.py deleted file mode 100644 index c0f19b5..0000000 --- a/tests/__main__.py +++ /dev/null @@ -1,6 +0,0 @@ -import unittest - -from test_basic import * -from test_advanced import * - -unittest.main() diff --git a/tests/context.py b/tests/context.py index bb675d4..fd3aeb0 100644 --- a/tests/context.py +++ b/tests/context.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - import sys, os sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) diff --git a/tests/test_advanced.py b/tests/test_advanced.py index a0bd976..3f6c382 100644 --- a/tests/test_advanced.py +++ b/tests/test_advanced.py @@ -1,7 +1,5 @@ import unittest -from context import pointcnn - class AdvancedTests(unittest.TestCase): """ Basic test cases """ diff --git a/tests/test_basic.py b/tests/test_basic.py index 47542ee..8793800 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,12 +1,30 @@ import unittest -from context import pointcnn +import torch +from torch.autograd import Variable +import numpy as np + +from pointcnn.core import XConv class BasicTests(unittest.TestCase): """ Basic test cases """ - def test_example(self): + def test_xconv_shape(self): self.assertTrue(True) + np.random.seed(0) + + N = 4 + D = 3 + C_in = 8 + 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)) if __name__ == "__main__": unittest.main()