wassname/PointCNN
1
0
Fork 0
mirror of https://github.com/wassname/PointCNN.git synced 2026-09-09 11:15:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

Vectorize XConv. Tests still required.

Browse Source
This commit is contained in:
Austin J. Garrett
2018-03-25 18:32:18 -04:00
parent a0f3d762fe
commit 9be3f8528a
3 changed files with 122 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pointcnn/context.py
+1 -1
View File
@@ -1,4 +1,4 @@
import sys, os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from lib import timed
# from lib import timed
pointcnn/core.py
+55 -100
View File
@@ -4,71 +4,21 @@ import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import matplotlib.pyplot as plt
try:
from .util import knn_indices_func
from .context import timed
from .util import knn_indices_func, MLP, BatchNorm, endchannels
# from .context import timed
except SystemError:
from util import knn_indices_func
from context import timed
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 appropriate axis into the channel dim, then tranpsose out.
"""
def __init__(self, D, num_features, dim = 1, *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.dim = dim
def forward(self, x):
x = torch.transpose(x, 1, self.dim).contiguous() # Must be made contiguous for cudNN.
self.bn(x)
x = torch.transpose(x, self.dim, 1)
return x
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, dim = -1, 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:])
])
from util import knn_indices_func, MLP, BatchNorm, endchannels
# from context import timed
class XConv(nn.Module):
"""
Vectorized pointwise convolution.
"""
def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None, mlp_width = 4):
def __init__(self, C_in, C_out, D, N_neighbors, N_rep, C_lifted = None, mlp_width = 2):
"""
:param C_in: Input dimension of the points' features.
:param C_out: Output dimension of the representative point features.
@@ -96,10 +46,9 @@ class XConv(nn.Module):
# Main dense linear layers
self.mlp_lift = MLP(np.around(np.geomspace(D, self.C_lifted, num = mlp_width)).astype(int))
self.mlp = nn.Sequential(
# torch.Conv2d(TODO),
MLP(np.around(np.geomspace(3, N_neighbors)).astype(int)) # Somehow, original code has K x K.
)
self.mid_conv = endchannels(nn.Conv2d(D, N_neighbors, 1))
self.mlp = MLP(np.around(np.geomspace(N_neighbors, N_neighbors)).astype(int)) # 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))
# Params for kernel initialization.
self.K = nn.Parameter(torch.FloatTensor(C_out, C_in + self.C_lifted, N_neighbors))
@@ -126,23 +75,20 @@ class XConv(nn.Module):
N = len(P)
p_center = torch.unsqueeze(p, dim = 2)
P_local = self.pts_batchnorm(P - p_center) # Move P to local coordinate system of p.
F_lifted = self.mlp_lift(P_local) # Individually lift each point into C_lifted dim space.
F_cat = torch.cat((F_lifted, F), -1) # Cat F_lifted and F, to size (N, N_rep, N_neighbors, C_lifted + C_in).
P_local = self.pts_batchnorm(P - p_center) # Move P to local coordinate system of p.
F_lifted = self.mlp_lift(P_local) # Individually lift each point into C_lifted dim space.
F_cat = torch.cat((F_lifted, F), -1) # Cat F_lifted and F, to size (N, N_rep, N_neighbors, C_lifted + C_in).
X_shape = (N, self.N_rep, N_neighbors, N_neighbors)
X = self.mlp(P_local).contiguous().view(*X_shape) # Learn the (N, K, K) X-transformation matrix.
F_X = torch.matmul(X, F_cat) # Weight and permute F_cat with the learned X.
# CODE PAST THIS POINT BROKEN.
# TODO: Implement separable_conv2d
F_p = nn.functional.conv1d( # Finally, typical convolution between K and F_X.
torch.transpose(F_X, 1, 2),
self.K
)
X = self.mlp(self.mid_conv(P_local)) # Learn the (N, K, K) X-transformation matrix.
X = X.contiguous().view(*X_shape)
F_X = torch.matmul(X, F_cat) # Weight and permute F_cat with the learned X.
F_p = self.end_conv(F_X)
return torch.squeeze(F_p, dim = 2)
class PointCNN(nn.Module):
"""
TODO: Insert documentation
"""
def __init__(self, C_in, C_out, D, N_neighbors, N_rep, r_indices_func, C_lifted = None, mlp_width = 4):
"""
@@ -151,7 +97,7 @@ class PointCNN(nn.Module):
:param D: Spatial dimensionality of points.
:param N_neighbors: Number of neighbors to convolve over.
:param r_indices_func: Selector function of the type,
INP
INP
======
ps : (N, N_rep, D) Representative points
P : (N, *, D) Point cloud
@@ -163,7 +109,7 @@ class PointCNN(nn.Module):
P[P_idx] is the set of points in the "region" around p.
a representative point p and a point cloud P. From these it returns an
array of N_neighbors
array of N_neighbors
:param C_lifted: Dimensionality of lifted point features.
:param mlp_width: Number of hidden layers in MLPs.
"""
@@ -177,19 +123,19 @@ class PointCNN(nn.Module):
def select_region(self, P, P_idx):
"""
Selects
Selects
:type P: FloatTensor (N, *, D)
:type P_idx: FloatTensor (N, N_neighbors)
:type P_idx: FloatTensor (N, N_rep, N_neighbors)
:rtype P_region: FloatTensor (N_rep, N_neighbors, D)
:param P: Point cloud to select regional points from
:param P_idx: Indices of points in region to be selected
:return:
"""
regions = torch.stack([
P[n][idx,:] for n, idx in enumerate(torch.unbind(P_idx, dim = 0))
], dim = 0)
return regions
@timed.timed
def forward(self, ps, P, F):
"""
Given a set of representative points, a point cloud, and its
@@ -204,50 +150,59 @@ class PointCNN(nn.Module):
:param F: Regional features such that P[:,p_idx,:] is the feature associated with F[:,p_idx,:]
:return:
"""
P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors).cuda()
inp_regions = torch.stack([
self.x_conv(p, self.select_region(P, P_idx[:,n]), self.select_region(F, P_idx[:,n]))
for n, p in enumerate(torch.unbind(ps, dim = 1))
], dim = 1)
return inp_regions
P_idx = self.r_indices_func(ps.cpu(), P.cpu(), N_neighbors)
P_regional = self.select_region(P, P_idx)
if False:
# Draw neighborhood points, for debugging.
t = 23
n = 3
test_point = ps[n,t,:].data.numpy()
neighborhood = P_regional[n,t,:,:].data.numpy()
plt.scatter(P[n][:,0], P[n][:,1])
plt.scatter(test_point[0], test_point[1], s = 100, c = 'green')
plt.scatter(neighborhood[:,0], neighborhood[:,1], s = 100, c = 'red')
plt.show()
F_regional = self.select_region(F, P_idx)
return self.x_conv(ps, P_regional, F_regional)
if __name__ == "__main__":
np.random.seed(0)
TESTING = XConv
TESTING = PointCNN
if TESTING == XConv:
N = 4
N_rep = 100
N_rep = 150
D = 3
C_in = 8
C_out = 32
N_neighbors = 10
model = XConv(C_in, C_out, D, N_neighbors, N_rep).cuda()
p = Variable(torch.from_numpy(np.random.rand(N,N_rep,D).astype(np.float32))).cuda()
P = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,D).astype(np.float32))).cuda()
F = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,C_in).astype(np.float32))).cuda()
model = XConv(C_in, C_out, D, N_neighbors, N_rep)
p = Variable(torch.from_numpy(np.random.rand(N,N_rep,D).astype(np.float32)))
P = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,D).astype(np.float32)))
F = Variable(torch.from_numpy(np.random.rand(N,N_rep,N_neighbors,C_in).astype(np.float32)))
out = model(p, P, F)
elif TESTING == PointCNN:
N = 4
num_points = 15000
N_rep = 7500
D = 3
C_in = 8
C_out = 32
N_neighbors = 5
num_points = 1000
N_rep = 50
D = 2
C_in = 64
C_out = 128
N_neighbors = 10
model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, knn_indices_func).cuda()
model = PointCNN(C_in, C_out, D, N_neighbors, N_rep, knn_indices_func)
test_P = np.random.rand(N,num_points,D).astype(np.float32)
test_F = np.random.rand(N,num_points,C_in).astype(np.float32)
idx = np.random.choice(test_P.shape[1], N_rep, replace = False)
test_ps = test_P[:,idx,:]
test_P = Variable(torch.from_numpy(test_P)).cuda()
test_F = Variable(torch.from_numpy(test_F)).cuda()
test_ps = Variable(torch.from_numpy(test_ps)).cuda()
test_P = Variable(torch.from_numpy(test_P))
test_F = Variable(torch.from_numpy(test_F))
test_ps = Variable(torch.from_numpy(test_ps))
out = model(test_ps, test_P, test_F)
pointcnn/util.py
+66 -1
View File
@@ -1,7 +1,62 @@
import torch
import torch.nn as nn
import numpy as np
from sklearn.neighbors import NearestNeighbors
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 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 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.
@@ -52,4 +107,14 @@ def knn_indices_func(ps, P, k):
if __name__ == "__main__":
from torch.autograd import Variable
N_rep = 100
N = 2
num_points = 1000
D = 3
test_P = np.random.rand(N,num_points,D).astype(np.float32)
idx = np.random.choice(test_P.shape[1], N_rep, replace = False)
test_ps = test_P[:,idx,:]
test_P = Variable(torch.from_numpy(test_P))
test_ps = Variable(torch.from_numpy(test_ps))
out = knn_indices_func(test_ps, test_P, 5)
print(out)