From 56e11ccd1e304fe1b80d45b91f5a438e5e542c1d Mon Sep 17 00:00:00 2001 From: Austin Garrett Date: Fri, 30 Mar 2018 17:59:17 -0400 Subject: [PATCH] Fix many parts of the model, most notably the convolutions. --- pointcnn/core.py | 79 ++++++++++++++++++++++++++++++++++------------ pointcnn/layers.py | 58 +++++++++++++++++----------------- pointcnn/util.py | 8 ----- 3 files changed, 87 insertions(+), 58 deletions(-) diff --git a/pointcnn/core.py b/pointcnn/core.py index f15cf39..b3013be 100644 --- a/pointcnn/core.py +++ b/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()) diff --git a/pointcnn/layers.py b/pointcnn/layers.py index 2900695..6a453bb 100644 --- a/pointcnn/layers.py +++ b/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: diff --git a/pointcnn/util.py b/pointcnn/util.py index 01c538a..623493c 100644 --- a/pointcnn/util.py +++ b/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.