diff --git a/pointcnn/core.py b/pointcnn/core.py index 83c7028..d852f54 100644 --- a/pointcnn/core.py +++ b/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()) diff --git a/pointcnn/layers.py b/pointcnn/layers.py new file mode 100644 index 0000000..2900695 --- /dev/null +++ b/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) diff --git a/pointcnn/util.py b/pointcnn/util.py index e29eb5f..f31a30c 100644 --- a/pointcnn/util.py +++ b/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.