Significantly clean and document code base. Project approximately complete.

core/__init__.py

@@ -0,0 +1,2 @@
+from PointCNN.core.model import XConv, PointCNN, RandPointCNN
+from PointCNN.core.util_funcs import knn_indices_func_cpu, knn_indices_func_gpu, UFloatTensor, ULongTensor

core/context.py

@@ -0,0 +1,5 @@
+import sys, os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# from lib import timed
+# from lib import pytorch_knn_cuda

core/model.py

@@ -0,0 +1,243 @@
+"""
+Author: Austin J. Garrett
+
+PyTorch implementation of the PointCNN paper, as specified in:
+  https://arxiv.org/pdf/1801.07791.pdf
+Original paper by: Yangyan Li, Rui Bu, Mingchao Sun, Baoquan Chen
+"""
+
+# External Modules
+import torch
+import torch.nn as nn
+from torch import FloatTensor
+import numpy as np
+from typing import Tuple, Callable, Optional
+
+# Internal Modules
+from PointCNN.core.util_funcs import UFloatTensor, ULongTensor
+from PointCNN.core.util_layers import Conv, SepConv, Dense, EndChannels
+
+class XConv(nn.Module):
+    """ Convolution over a single point and its neighbors.  """
+
+    def __init__(self, C_in : int, C_out : int, dims : int, K : int,
+                 P : int, C_mid : int, depth_multiplier : int) -> None:
+        """
+        :param C_in: Input dimension of the points' features.
+        :param C_out: Output dimension of the representative point features.
+        :param dims: Spatial dimensionality of points.
+        :param K: Number of neighbors to convolve over.
+        :param P: Number of representative points.
+        :param C_mid: Dimensionality of lifted point features.
+        :param depth_multiplier: Depth multiplier for internal depthwise separable convolution.
+        """
+        super(XConv, self).__init__()
+
+        if __debug__:
+            # Only needed for assertions.
+            self.C_in = C_in
+            self.C_mid = C_mid
+            self.dims = dims
+            self.K = K
+
+        self.P = P
+
+        # Additional processing layers
+        # self.pts_layernorm = LayerNorm(2, momentum = 0.9)
+
+        # Main dense linear layers
+        self.dense1 = Dense(dims, C_mid)
+        self.dense2 = Dense(C_mid, C_mid)
+
+        # Layers to generate X
+        self.x_trans = nn.Sequential(
+            EndChannels(Conv(
+                in_channels = dims,
+                out_channels = K*K,
+                kernel_size = (1, K),
+                with_bn = False
+            )),
+            Dense(K*K, K*K, with_bn = False),
+            Dense(K*K, K*K, with_bn = False, activation = None)
+        )
+
+        self.end_conv = EndChannels(SepConv(
+            in_channels = C_mid + C_in,
+            out_channels = C_out,
+            kernel_size = (1, K),
+            depth_multiplier = depth_multiplier
+        )).cuda()
+
+    def forward(self, x : Tuple[UFloatTensor,            # (N, P, dims)
+                                UFloatTensor,            # (N, P, K, dims)
+                                Optional[UFloatTensor]]  # (N, P, K, C_in)
+               ) -> UFloatTensor:                        # (N, K, C_out)
+        """
+        Applies XConv to the input data.
+        :param x: (rep_pt, pts, fts) where
+          - rep_pt: Representative point.
+          - pts: Regional point cloud such that fts[:,p_idx,:] is the feature
+          associated with pts[:,p_idx,:].
+          - fts: Regional features such that pts[:,p_idx,:] is the feature
+          associated with fts[:,p_idx,:].
+        :return: Features aggregated into point rep_pt.
+        """
+        rep_pt, pts, fts = x
+
+        if fts is not None:
+            assert(rep_pt.size()[0] == pts.size()[0] == fts.size()[0])  # Check N is equal.
+            assert(rep_pt.size()[1] == pts.size()[1] == fts.size()[1])  # Check P is equal.
+            assert(pts.size()[2] == fts.size()[2] == self.K)            # Check K is equal.
+            assert(fts.size()[3] == self.C_in)                          # Check C_in is equal.
+        else:
+            assert(rep_pt.size()[0] == pts.size()[0])                   # Check N is equal.
+            assert(rep_pt.size()[1] == pts.size()[1])                   # Check P is equal.
+            assert(pts.size()[2] == self.K)                             # Check K is equal.
+        assert(rep_pt.size()[2] == pts.size()[3] == self.dims)          # Check dims is equal.
+
+        N = len(pts)
+        P = rep_pt.size()[1]  # (N, P, K, dims)
+        p_center = torch.unsqueeze(rep_pt, dim = 2)  # (N, P, 1, dims)
+
+        # Move pts to local coordinate system of rep_pt.
+        pts_local = pts - p_center  # (N, P, K, dims)
+        # pts_local = self.pts_layernorm(pts - p_center)
+
+        # Individually lift each point into C_mid space.
+        fts_lifted0 = self.dense1(pts_local)
+        fts_lifted  = self.dense2(fts_lifted0)  # (N, P, K, C_mid)
+
+        if fts is None:
+            fts_cat = fts_lifted
+        else:
+            fts_cat = torch.cat((fts_lifted, fts), -1)  # (N, P, K, C_mid + C_in)
+
+        # Learn the (N, K, K) X-transformation matrix.
+        X_shape = (N, P, self.K, self.K)
+        X = self.x_trans(pts_local)
+        X = X.view(*X_shape)
+
+        # Weight and permute fts_cat with the learned X.
+        fts_X = torch.matmul(X, fts_cat)
+        fts_p = self.end_conv(fts_X).squeeze(dim = 2)
+        return fts_p
+
+class PointCNN(nn.Module):
+    """ Pointwise convolutional model. """
+
+    def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
+                 r_indices_func : Callable[[UFloatTensor,  # (N, P, dims)
+                                            UFloatTensor,  # (N, x, dims)
+                                            int, int],
+                                           ULongTensor]    # (N, P, K)
+                ) -> None:
+        """
+        :param C_in: Input dimension of the points' features.
+        :param C_out: Output dimension of the representative point features.
+        :param dims: Spatial dimensionality of points.
+        :param K: Number of neighbors to convolve over.
+        :param D: "Spread" of neighboring points.
+        :param P: Number of representative points.
+        :param r_indices_func: Selector function of the type,
+          INPUTS
+          rep_pts : Representative points.
+          pts  : Point cloud.
+          K : Number of points for each region.
+          D : "Spread" of neighboring points.
+
+          OUTPUT
+          pts_idx : Array of indices into pts such that pts[pts_idx] is the set
+          of points in the "region" around rep_pt.
+        """
+        super(PointCNN, self).__init__()
+
+        C_mid = C_out // 2 if C_in == 0 else C_out // 4
+        depth_multiplier = min(int(np.ceil(C_out / C_in)), 4)
+
+        self.r_indices_func = lambda rep_pts, pts: r_indices_func(rep_pts, pts, K, D)
+        self.dense = Dense(C_in, C_out // 2) if C_in != 0 else None
+        self.x_conv = XConv(C_out // 2 if C_in != 0 else C_in, C_out, dims, K, P, C_mid, depth_multiplier)
+        self.D = D
+
+    def select_region(self, pts : UFloatTensor,  # (N, x, dims)
+                      pts_idx : ULongTensor      # (N, P, K)
+                     ) -> UFloatTensor:          # (P, K, dims)
+        """
+        Selects neighborhood points based on output of r_indices_func.
+        :param pts: Point cloud to select regional points from.
+        :param pts_idx: Indices of points in region to be selected.
+        :return: Local neighborhoods around each representative point.
+        """
+        regions = torch.stack([
+            pts[n][idx,:] for n, idx in enumerate(torch.unbind(pts_idx, dim = 0))
+        ], dim = 0)
+        return regions
+
+    def forward(self, x : Tuple[FloatTensor,  # (N, P, dims)
+                                FloatTensor,  # (N, x, dims)
+                                FloatTensor]  # (N, x, C_in)
+               ) -> FloatTensor:              # (N, P, C_out)
+        """
+        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.
+        :param x: (rep_pts, pts, fts) where
+          - rep_pts: Representative points.
+          - pts: Regional point cloud such that fts[:,p_idx,:] is the
+          feature associated with pts[:,p_idx,:].
+          - fts: Regional features such that pts[:,p_idx,:] is the feature
+          associated with fts[:,p_idx,:].
+        :return: Features aggregated to rep_pts.
+        """
+        rep_pts, pts, fts = x
+        fts = self.dense(fts) if fts is not None else fts
+
+        # This step takes ~97% of the time. Prime target for optimization: KNN on GPU.
+        pts_idx = self.r_indices_func(rep_pts.cpu(), pts.cpu()).cuda()
+        # -------------------------------------------------------------------------- #
+
+        pts_regional = self.select_region(pts, pts_idx)
+        fts_regional = self.select_region(fts, pts_idx) if fts is not None else fts
+        fts_p = self.x_conv((rep_pts, pts_regional, fts_regional))
+
+        return fts_p
+
+class RandPointCNN(nn.Module):
+    """ PointCNN with randomly subsampled representative points. """
+
+    def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
+                 r_indices_func : Callable[[UFloatTensor,  # (N, P, dims)
+                                            UFloatTensor,  # (N, x, dims)
+                                            int, int],
+                                           ULongTensor]    # (N, P, K)
+                ) -> None:
+        """ See documentation for PointCNN. """
+        super(RandPointCNN, self).__init__()
+        self.pointcnn = PointCNN(C_in, C_out, dims, K, D, P, r_indices_func)
+        self.P = P
+
+    def forward(self, x : Tuple[UFloatTensor,  # (N, x, dims)
+                                UFloatTensor]  # (N, x, dims)
+               ) -> Tuple[UFloatTensor,        # (N, P, dims)
+                          UFloatTensor]:       # (N, P, C_out)
+        """
+        Given a point cloud, and its corresponding features, return a new set
+        of randomly-sampled representative points with features projected from
+        the point cloud.
+        :param x: (pts, fts) where
+         - pts: Regional point cloud such that fts[:,p_idx,:] is the
+        feature associated with pts[:,p_idx,:].
+         - fts: Regional features such that pts[:,p_idx,:] is the feature
+        associated with fts[:,p_idx,:].
+        :return: Randomly subsampled points and their features.
+        """
+        pts, fts = x
+        if 0 < self.P < pts.size()[1]:
+            # Select random set of indices of subsampled points.
+            idx = np.random.choice(pts.size()[1], self.P, replace = False).tolist()
+            rep_pts = pts[:,idx,:]
+        else:
+            # All input points are representative points.
+            rep_pts = pts
+        rep_pts_fts = self.pointcnn((rep_pts, pts, fts))
+        return rep_pts, rep_pts_fts

core/util_funcs.py

@@ -0,0 +1,66 @@
+# External Modules
+import torch
+from torch import cuda, FloatTensor, LongTensor
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.neighbors import NearestNeighbors
+from typing import Union
+
+# Types to allow for both CPU and GPU models.
+UFloatTensor = Union[FloatTensor, cuda.FloatTensor]
+ULongTensor = Union[LongTensor, cuda.LongTensor]
+
+def knn_indices_func_cpu(rep_pts : FloatTensor,  # (N, pts, dim)
+                         pts : FloatTensor,      # (N, x, dim)
+                         K : int, D : int
+                        ) -> LongTensor:         # (N, pts, K)
+    """
+    CPU-based Indexing function based on K-Nearest Neighbors search.
+    :param rep_pts: Representative points.
+    :param pts: Point cloud to get indices from.
+    :param K: Number of nearest neighbors to collect.
+    :param D: "Spread" of neighboring points.
+    :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:]
+    is the set k-nearest neighbors for the representative points in pts[n].
+    """
+    rep_pts = rep_pts.data.numpy()
+    pts = pts.data.numpy()
+    region_idx = []
+
+    for n, p in enumerate(rep_pts):
+        P_particular = pts[n]
+        nbrs = NearestNeighbors(D*K + 1, algorithm = "ball_tree").fit(P_particular)
+        indices = nbrs.kneighbors(p)[1]
+        region_idx.append(indices[:,1::D])
+
+    region_idx = torch.from_numpy(np.stack(region_idx, axis = 0))
+    return region_idx
+
+def knn_indices_func_gpu(rep_pts : cuda.FloatTensor,  # (N, pts, dim)
+                         pts : cuda.FloatTensor,      # (N, x, dim)
+                         k : int, d : int
+                        ) -> cuda.LongTensor:         # (N, pts, K)
+    """
+    GPU-based Indexing function based on K-Nearest Neighbors search.
+    Very memory intensive, and thus unoptimal for large numbers of points.
+    :param rep_pts: Representative points.
+    :param pts: Point cloud to get indices from.
+    :param K: Number of nearest neighbors to collect.
+    :param D: "Spread" of neighboring points.
+    :return: Array of indices, P_idx, into pts such that pts[n][P_idx[n],:]
+    is the set k-nearest neighbors for the representative points in pts[n].
+    """
+    region_idx = []
+
+    for n, qry in enumerate(rep_pts):
+        ref = pts[n]
+        n, d = ref.size()
+        m, d = qry.size()
+        mref = ref.expand(m, n, d)
+        mqry = qry.expand(n, m, d).transpose(0, 1)
+        dist2 = torch.sum((mqry - mref)**2, 2).squeeze()
+        _, inds = torch.topk(dist2, k*d + 1, dim = 1, largest = False)
+        region_idx.append(inds[:,1::d])
+
+    region_idx = torch.stack(region_idx, dim = 0)
+    return region_idx

core/util_layers.py

@@ -0,0 +1,148 @@
+import torch.nn as nn
+from typing import Callable, Union, Tuple
+
+from PointCNN.core.util_funcs import UFloatTensor
+
+def EndChannels(f, make_contiguous = False):
+    """ Class decorator to apply 2D convolution along end channels. """
+
+    class WrappedLayer(nn.Module):
+
+        def __init__(self):
+            super(WrappedLayer, self).__init__()
+            self.forward = lambda x: f(x.permute(0,3,1,2)).permute(0,2,3,1)
+
+    return WrappedLayer()
+
+class Dense(nn.Module):
+    """
+    Single layer perceptron with optional activation, batch normalization, and dropout.
+    """
+
+    def __init__(self, in_features : int, out_features : int,
+                 drop_rate : int = 0, with_bn : bool = True,
+                 activation : Callable[[UFloatTensor], UFloatTensor] = nn.ReLU()
+                ) -> None:
+        """
+        :param in_features: Length of input featuers (last dimension).
+        :param out_features: Length of output features (last dimension).
+        :param drop_rate: Drop rate to be applied after activation.
+        :param with_bn: Whether or not to apply batch normalization.
+        :param activation: Activation function.
+        """
+        super(Dense, self).__init__()
+
+        self.linear = nn.Linear(in_features, out_features)
+        self.activation = activation
+        # self.bn = LayerNorm(out_channels) if with_bn else None
+        self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None
+
+    def forward(self, x : UFloatTensor) -> UFloatTensor:
+        """
+        :param x: Any input tensor that can be input into nn.Linear.
+        :return: Tensor with linear layer and optional activation, batchnorm,
+        and dropout applied.
+        """
+        x = self.linear(x)
+        if self.activation:
+            x = self.activation(x)
+        # if self.bn:
+        #     x = self.bn(x)
+        if self.drop:
+            x = self.drop(x)
+        return x
+
+class Conv(nn.Module):
+    """
+    2D convolutional layer with optional activation and batch normalization.
+    """
+
+    def __init__(self, in_channels : int, out_channels : int,
+                 kernel_size : Union[int, Tuple[int, int]], with_bn : bool = True,
+                 activation : Callable[[UFloatTensor], UFloatTensor] = nn.relu()
+                ) -> None:
+        """
+        :param in_channels: Length of input featuers (first dimension).
+        :param out_channels: Length of output features (first dimension).
+        :param kernel_size: Size of convolutional kernel.
+        :param with_bn: Whether or not to apply batch normalization.
+        :param activation: Activation function.
+        """
+        super(Conv, self).__init__()
+
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias = not with_bn)
+        self.activation = activation
+        self.bn = nn.BatchNorm2d(out_channels, momentum = 0.9) if with_bn else None
+
+    def forward(self, x : UFloatTensor) -> UFloatTensor:
+        """
+        :param x: Any input tensor that can be input into nn.Conv2d.
+        :return: Tensor with convolutional layer and optional activation and batchnorm applied.
+        """
+        x = self.conv(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.bn:
+            x = self.bn(x)
+        return x
+
+class SepConv(nn.Module):
+    """ Depthwise separable convolution with optional activation and batch normalization"""
+
+    def __init__(self, in_channels : int, out_channels : int,
+                 kernel_size : Union[int, Tuple[int, int]],
+                 depth_multiplier : int = 1, with_bn : bool = True,
+                 activation : Callable[[UFloatTensor], UFloatTensor] = nn.ReLU()
+                 ) -> None:
+        """
+        :param in_channels: Length of input featuers (first dimension).
+        :param out_channels: Length of output features (first dimension).
+        :param kernel_size: Size of convolutional kernel.
+        :depth_multiplier: Depth multiplier for middle part of separable convolution.
+        :param with_bn: Whether or not to apply batch normalization.
+        :param activation: Activation function.
+        """
+        super(SepConv, self).__init__()
+
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels * depth_multiplier, kernel_size, groups = in_channels),
+            nn.Conv2d(in_channels * depth_multiplier, out_channels, 1, bias = not with_bn)
+        )
+
+        self.activation = activation
+        self.bn = nn.BatchNorm2d(out_channels, momentum = 0.9) if with_bn else None
+
+    def forward(self, x : UFloatTensor) -> UFloatTensor:
+        """
+        :param x: Any input tensor that can be input into nn.Conv2d.
+        :return: Tensor with depthwise separable convolutional layer and
+        optional activation and batchnorm applied.
+        """
+        x = self.conv(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.bn:
+            x = self.bn(x)
+        return x
+
+class LayerNorm(nn.Module):
+    """
+    Batch Normalization over ONLY the mini-batch layer (suitable for nn.Linear layers).
+    """
+
+    def __init__(self, N : int, dim : int, *args, **kwargs) -> None:
+        """
+        :param N: Batch size.
+        :param D: Dimensions.
+        """
+        super(LayerNorm, self).__init__()
+        if dim == 1:
+            self.bn = nn.BatchNorm1d(N, *args, **kwargs)
+        elif dim == 2:
+            self.bn = nn.BatchNorm2d(N, *args, **kwargs)
+        elif dim == 3:
+            self.bn = nn.BatchNorm3d(N, *args, **kwargs)
+        else:
+            raise ValueError("Dimensionality %i not supported" % dim)
+
+        self.forward = lambda x: self.bn(x.unsqueeze(0)).squeeze(0)