A bit of refactoring.

pointcnn/core.py

@@ -7,10 +7,12 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 try:
-    from .util import knn_indices_func, MLP, BatchNorm, endchannels
+    from .util import knn_indices_func, endchannels
+    from .layers import MLP, BatchNorm, SeparableConv2d, DepthwiseConv2d
     from .context import timed
 except SystemError:
-    from util import knn_indices_func, MLP, BatchNorm, endchannels
+    from util import knn_indices_func, endchannels
+    from layers import MLP, BatchNorm, SeparableConv2d, DepthwiseConv2d
     from context import timed
 
 class XConv(nn.Module):
@@ -46,9 +48,9 @@ class XConv(nn.Module):
         # 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.mlp_lift = MLP([D] + [self.C_lifted] * mlp_width)
         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.mlp = MLP([N_neighbors] * mlp_width)  # 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).cuda())
 
         # Params for kernel initialization.
@@ -154,7 +156,7 @@ class PointCNN(nn.Module):
         """
         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:
+        if False:
             # Draw neighborhood points, for debugging.
             t = 23
             n = 3
@@ -165,6 +167,7 @@ class PointCNN(nn.Module):
             plt.scatter(neighborhood[:,0], neighborhood[:,1], s = 100, c = 'red')
             plt.show()
         F_regional = self.select_region(F, P_idx)
+        # ps, P, F_P -> ps_F
         return self.x_conv(ps, P_regional, F_regional)
 
 if __name__ == "__main__":
@@ -207,5 +210,6 @@ if __name__ == "__main__":
         test_F = Variable(torch.from_numpy(test_F)).cuda()
         test_ps = Variable(torch.from_numpy(test_ps)).cuda()
 
-        for _ in range(10):
-            out = model(test_ps, test_P, test_F)
+        print(test_F.size())
+        out = model(test_ps, test_P, test_F)
+        print(out.size())

pointcnn/layers.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+try:
+    from .util import endchannels
+except:
+    from util import endchannels
+
+def SeparableConv2d(in_channels, out_channels, kernel_size):
+    """
+    Separable convolution (is this correct?)
+    """
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, kernel_size = (1, kernel_size)),
+        nn.Conv2d(out_channels, in_channels, kernel_size = (kernel_size, 1), groups = in_channels)
+    )
+
+def DepthwiseConv2d(in_channel, depth_multiplier):
+    """
+    Factory function to generate depthwise 2d convolutional layer.
+    :param in_channel: TODO
+    :param out_channel: TODO
+    :param depth_multiplier: TODO
+    :return: TODO
+    """
+    return nn.Conv2d(in_channels, depth_multiplier * in_channels, groups = in_channels)
+
+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 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:])
+        ])
+
+if __name__ == "__main__":
+    ftr_map = torch.autograd.Variable(torch.FloatTensor(2,8,100,100))
+    layer = SeparableConv2d(8, 16, 2)
+    out = layer(ftr_map)
+    print(out)
+
+    test = nn.SpatialConvolutionLocal(8, 16, 100, 100, 100, 100)

pointcnn/util.py

@@ -1,7 +1,4 @@
-import time
-
 import torch
-import torch.nn as nn
 
 import numpy as np
 from sklearn.neighbors import NearestNeighbors
@@ -19,51 +16,6 @@ def endchannels(f, make_contiguous = False):
             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) ->
-    (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 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.