Implement XConv operator. Begin implementing PointCNN.

Makefile

@@ -2,4 +2,4 @@ init:
 	pip3 install -r requirements.txt
 
 test:
-	nosetests tests
+	python3.5 -m unittest discover

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

pointcnn/pointcnn.py

@@ -1 +0,0 @@
-# This is a file for the main PyTorch module.

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)

tests/__main__.py

@@ -1,6 +0,0 @@
-import unittest
-
-from test_basic import *
-from test_advanced import *
-
-unittest.main()

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__), '..')))
 

tests/test_advanced.py

@@ -1,7 +1,5 @@
 import unittest
 
-from context import pointcnn
-
 class AdvancedTests(unittest.TestCase):
     """ Basic test cases """
 

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()