Fix many parts of the model, most notably the convolutions.

pointcnn/core.py

@@ -7,12 +7,12 @@ import numpy as np
 import matplotlib.pyplot as plt
 
 try:
-    from .util import knn_indices_func, endchannels
-    from .layers import MLP, BatchNorm, SeparableConv2d, DepthwiseConv2d
+    from .util import knn_indices_func
+    from .layers import MLP, LayerNorm, DepthwiseSeparableConv2d, endchannels
     from .context import timed
 except SystemError:
-    from util import knn_indices_func, endchannels
-    from layers import MLP, BatchNorm, SeparableConv2d, DepthwiseConv2d
+    from util import knn_indices_func
+    from layers import MLP, LayerNorm, DepthwiseSeparableConv2d, endchannels
     from context import timed
 
 class XConv(nn.Module):
@@ -44,20 +44,53 @@ class XConv(nn.Module):
         self.N_rep = N_rep
 
         # Additional processing layers
-        self.pts_batchnorm = BatchNorm(2, D, momentum = 0.9)
-        # self.pts_batchnorm = BatchNorm(BatchNorm())
+        # self.pts_layernorm = LayerNorm(2, momentum = 0.9)
 
         # Main dense linear layers
-        self.mlp_lift = MLP([D] + [self.C_lifted] * mlp_width)
-        self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1).cuda())
-        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)).cuda())
+        self.mlp_lift = MLP([D] + [self.C_lifted] * (mlp_width - 1), batch_norm = False)
+        # self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1).cuda())
 
-        # 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)
+        # Layers to generate X
+        self.mid_conv = endchannels(nn.Sequential(
+            nn.Conv2d(D, N_neighbors**2, (1, N_neighbors)).cuda(),
+            nn.ReLU(),
+            nn.BatchNorm2d(N_neighbors**2).cuda(),
+        ))
+        self.mid_dwconv1 = endchannels(nn.Sequential(
+            DepthwiseSeparableConv2d(
+                in_channels = N_neighbors,
+                out_channels = N_neighbors**2,
+                kernel_size = (1, N_neighbors),
+                depth_multiplier = N_neighbors
+            ).cuda(),
+            nn.ReLU(),
+            nn.BatchNorm2d(N_neighbors ** 2).cuda()
+        ))
+        self.mid_dwconv2 = endchannels(nn.Sequential(
+            DepthwiseSeparableConv2d(
+                in_channels = N_neighbors,
+                out_channels = N_neighbors**2,
+                kernel_size = (1, N_neighbors),
+                depth_multiplier = N_neighbors
+            ).cuda(),
+            nn.ReLU(),
+            nn.BatchNorm2d(N_neighbors**2).cuda()
+        ))
 
+        # Final 
+        self.mlp = MLP([N_neighbors] * mlp_width, batch_norm = False)
+        self.end_conv = endchannels(nn.Sequential(
+            DepthwiseSeparableConv2d(
+                in_channels = C_lifted + C_in,
+                out_channels = C_out,
+                kernel_size = (1, N_neighbors),
+                depth_multiplier = 4
+            ).cuda(),
+            nn.ReLU(),
+            nn.BatchNorm2d(C_out).cuda()
+        ))
+
+    # @timed.timed
     def forward(self, x):
         """
         Applies XConv to the input data.
@@ -82,7 +115,8 @@ class XConv(nn.Module):
         p_center = torch.unsqueeze(p, dim = 2)
 
         # Move P to local coordinate system of p.
-        P_local = self.pts_batchnorm(P - p_center)
+        P_local = P - p_center
+        # P_local = self.pts_layernorm(P - p_center)
 
         # Individually lift each point into C_lifted dim space.
         F_lifted = self.mlp_lift(P_local)
@@ -92,14 +126,17 @@ class XConv(nn.Module):
 
         # Learn the (N, K, K) X-transformation matrix.
         X_shape = (N, N_rep, self.N_neighbors, self.N_neighbors)
-        X = self.mlp(self.mid_conv(P_local))
+        X = self.mid_conv(P_local)
+        X = X.contiguous().view(*X_shape)
+        X = self.mid_dwconv1(X)
+        X = X.contiguous().view(*X_shape)
+        X = self.mid_dwconv2(X)
         X = X.contiguous().view(*X_shape)
 
         # Weight and permute F_cat with the learned X.
         F_X = torch.matmul(X, F_cat)
-        F_p = self.end_conv(F_X)
-        time.sleep(5)
-        return torch.squeeze(F_p, dim = 2)
+        F_p = self.end_conv(F_X).squeeze(dim = 2)
+        return F_p
 
 class PointCNN(nn.Module):
     """
@@ -152,7 +189,6 @@ class PointCNN(nn.Module):
         ], dim = 0)
         return regions
 
-    # @timed.timed
     def forward(self, x):
         """
         Given a set of representative points, a point cloud, and its
@@ -225,5 +261,6 @@ if __name__ == "__main__":
         test_ps = Variable(torch.from_numpy(test_ps)).cuda()
 
         print(test_F.size())
-        out = model((test_ps, test_P, test_F))
+        for _ in range(50):
+            out = model((test_ps, test_P, test_F))
         print(out.size())

pointcnn/layers.py

@@ -3,50 +3,50 @@ import torch.nn as nn
 import numpy as np
 
 try:
-    from .util import endchannels
+    pass
+    # from .util import endchannels
 except:
-    from util import endchannels
+    # from util import endchannels
+    pass
 
-def SeparableConv2d(in_channels, out_channels, kernel_size):
+def endchannels(f, make_contiguous = False):
+    class wrapped_layer(nn.Module):
+        def __init__(self):
+            super(wrapped_layer, self).__init__()
+        def forward(self, x):
+            x = x.permute(0,3,1,2)
+            x = f(x)
+            x = x.permute(0,2,3,1)
+            return x
+    return wrapped_layer()
+
+def DepthwiseSeparableConv2d(in_channels, out_channels, kernel_size, depth_multiplier):
     """
-    Separable convolution (is this correct?)
+    Depthwise separable convolution
     """
     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)
+        nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels),
+        nn.Conv2d(in_channels * depth_multiplier, out_channels, 1)
     )
 
-def DepthwiseConv2d(in_channel, depth_multiplier):
+class LayerNorm(nn.Module):
     """
-    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.
+    Batch Normalization over ONLY the mini-batch layer
+    (suitable for nn.Linear layers).
     """
 
-    def __init__(self, D, num_features, *args, **kwargs):
-        super(BatchNorm, self).__init__()
+    def __init__(self, N, D, *args, **kwargs):
+        super(LayerNorm, self).__init__()
         if D == 1:
-            self.bn = nn.BatchNorm1d(num_features, *args, **kwargs)
+            self.bn = nn.BatchNorm1d(N, *args, **kwargs)
         elif D == 2:
-            self.bn = nn.BatchNorm2d(num_features, *args, **kwargs)
+            self.bn = nn.BatchNorm2d(N, *args, **kwargs)
         elif D == 3:
-            self.bn = nn.BatchNorm3d(num_features, *args, **kwargs)
+            self.bn = nn.BatchNorm3d(N, *args, **kwargs)
         else:
             raise ValueError("Dimensionality %i not supported" % D)
 
-        self.forward = endchannels(self.bn, make_contiguous = True)
+        self.forward = lambda x: self.bn(x.unsqueeze(0)).squeeze(0)
 
 def MLP(layer_sizes, activation_layer = nn.ReLU(), batch_norm = True):
     """
@@ -61,7 +61,7 @@ def MLP(layer_sizes, activation_layer = nn.ReLU(), batch_norm = True):
         return nn.Sequential(*[
             nn.Sequential(nn.Linear(C_in, C_out),
                           activation_layer,
-                          BatchNorm(D = 2, num_features = C_out, momentum = 0.9)
+                          LayerNorm(D = 2, momentum = 0.9)
             ) for (C_in, C_out) in zip(layer_sizes, layer_sizes[1:])
         ])
     else:

pointcnn/util.py

@@ -10,14 +10,6 @@ try:
 except SystemError:
     from context import pytorch_knn_cuda
 
-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 apply_along_dim(xs, f, dim):
     """
     PyTorch analog to np.apply_along_axis.