mirror of
https://github.com/wassname/PointCNN.git
synced 2026-09-08 16:50:50 +08:00
Significantly clean and document code base. Project approximately complete.
This commit is contained in:
Austin J. Garrett
@@ -1,2 +1,9 @@
|
||||
# PointCNN
|
||||
PyTorch implementation of PointCNN model specified in the white paper located here: https://arxiv.org/pdf/1801.07791.pdf
|
||||
|
||||
Current MNIST accuracy: ~96%
|
||||
|
||||
My coding style is somewhat unique, but ultimately geared towards maximal
|
||||
readability. Along with extensive documentation in the code, I use type , and
|
||||
code comments indicating input/outputs shapes. (x,y,z) just indicate that any
|
||||
value is accepted at runtime.
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from PointCNN.core import RandPointCNN, PointCNN, XConv, knn_indices_func_cpu, knn_indices_func_gpu
|
||||
@@ -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
|
||||
@@ -1,40 +1,24 @@
|
||||
"""
|
||||
PyTorch implementation of the PointCNN paper, as specified in:
|
||||
https://arxiv.org/pdf/1801.07791.pdf
|
||||
|
||||
Author: Austin J. Garrett
|
||||
|
||||
I make liberal use of the mypy static type checker for Python.
|
||||
It should be mostly intuitive, but further documentation can be found at:
|
||||
http://mypy-lang.org/
|
||||
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
|
||||
"""
|
||||
|
||||
# Standard Modules
|
||||
import time
|
||||
|
||||
# External Modules
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor, LongTensor
|
||||
from torch.autograd import Variable
|
||||
from torch import FloatTensor
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from typing import Tuple, Callable
|
||||
from typing import Tuple, Callable, Optional
|
||||
|
||||
# Internal Modules
|
||||
try:
|
||||
from .util import knn_indices_func, knn_indices_func_gpu, plot
|
||||
from .layers import MLP, LayerNorm, Conv, SepConv, Dense, end_channels
|
||||
# from .context import timed
|
||||
except SystemError:
|
||||
from util import knn_indices_func, knn_indices_func_gpu, plot
|
||||
from layers import MLP, LayerNorm, Conv, SepConv, Dense, end_channels
|
||||
# from context import timed
|
||||
from PointCNN.core.util_funcs import UFloatTensor, ULongTensor
|
||||
from PointCNN.core.util_layers import Conv, SepConv, Dense, EndChannels
|
||||
|
||||
class XConv(nn.Module):
|
||||
"""
|
||||
Vectorized pointwise convolution.
|
||||
"""
|
||||
""" 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:
|
||||
@@ -43,7 +27,7 @@ class XConv(nn.Module):
|
||||
: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 P: Number of representative points.
|
||||
:param C_mid: Dimensionality of lifted point features.
|
||||
:param depth_multiplier: Depth multiplier for internal depthwise separable convolution.
|
||||
"""
|
||||
@@ -67,7 +51,7 @@ class XConv(nn.Module):
|
||||
|
||||
# Layers to generate X
|
||||
self.x_trans = nn.Sequential(
|
||||
end_channels(Conv(
|
||||
EndChannels(Conv(
|
||||
in_channels = dims,
|
||||
out_channels = K*K,
|
||||
kernel_size = (1, K),
|
||||
@@ -77,30 +61,29 @@ class XConv(nn.Module):
|
||||
Dense(K*K, K*K, with_bn = False, activation = None)
|
||||
)
|
||||
|
||||
self.end_conv = end_channels(SepConv(
|
||||
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[Tensor, Tensor, Tensor]) -> Tensor:
|
||||
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.
|
||||
:type rep_pt: (N, P, dims)
|
||||
:type pts: (N, P, K, dims)
|
||||
:type fts: (N, P, K, C_in)
|
||||
:rtype: (TODO: shape)
|
||||
:param x: (rep_pt, pts, fts)
|
||||
:param rep_pt: Representative point
|
||||
:param pts: Regional point cloud such that fts[:,p_idx,:] is the feature associated with pts[:,p_idx,:]
|
||||
:param fts: Regional features such that pts[:,p_idx,:] is the feature associated with fts[:,p_idx,:]
|
||||
: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
|
||||
N = len(pts)
|
||||
|
||||
#== RUNTIME ASSERTIONS ==#
|
||||
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.
|
||||
@@ -111,15 +94,15 @@ class XConv(nn.Module):
|
||||
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)
|
||||
@@ -134,125 +117,120 @@ class XConv(nn.Module):
|
||||
X = self.x_trans(pts_local)
|
||||
X = X.view(*X_shape)
|
||||
|
||||
"""
|
||||
X = self.mid_conv(pts_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 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):
|
||||
"""
|
||||
TODO: Insert documentation
|
||||
"""
|
||||
""" Pointwise convolutional model. """
|
||||
|
||||
def __init__(self, C_in : int, C_out : int, dims : int, K : int, D : int, P : int,
|
||||
r_indices_func : Callable[[Tensor, Tensor, int, int], LongTensor]) -> None:
|
||||
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 P: Number of representative points.
|
||||
:param D: "Spread" of neighboring points.
|
||||
:param P: Number of representative points.
|
||||
:param r_indices_func: Selector function of the type,
|
||||
INP
|
||||
======
|
||||
rep_pts : (N, P, dims) Representative points
|
||||
pts : (N, *, dims) Point cloud
|
||||
K : Number of points for each region.
|
||||
D : "Spread" of neighboring points (analogous to stride).
|
||||
INPUTS
|
||||
rep_pts : Representative points.
|
||||
pts : Point cloud.
|
||||
K : Number of points for each region.
|
||||
D : "Spread" of neighboring points.
|
||||
|
||||
OUT
|
||||
======
|
||||
pts_idx : (N, P, K) Array of indices into pts such that
|
||||
pts[pts_idx] is the set of points in the "region" around rep_pt.
|
||||
|
||||
a representative point rep_pt and a point cloud pts. From these it returns an
|
||||
array of K
|
||||
:param C_mid: Dimensionality of lifted point features.
|
||||
:param mlp_width: Number of hidden layers in MLPs.
|
||||
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 = r_indices_func
|
||||
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 : Tensor, pts_idx : LongTensor) -> Tensor:
|
||||
def select_region(self, pts : UFloatTensor, # (N, x, dims)
|
||||
pts_idx : ULongTensor # (N, P, K)
|
||||
) -> UFloatTensor: # (P, K, dims)
|
||||
"""
|
||||
Selects
|
||||
:type pts: (N, *, dims)
|
||||
:type pts_idx: (N, P, K)
|
||||
:rtype pts_region: (P, K, dims)
|
||||
:param pts: Point cloud to select regional points from
|
||||
:param pts_idx: Indices of points in region to be selected
|
||||
:return:
|
||||
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[Tensor, Tensor, Tensor]) -> Tensor:
|
||||
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.
|
||||
:type rep_pts: (N, *, dims)
|
||||
:type pts: (N, K, dims)
|
||||
:type fts: (N, K, C_in)
|
||||
:rtype: (N, P, dims)
|
||||
:param rep_pts: Representative points
|
||||
:param pts: Regional point cloud such that fts[:,p_idx,:] is the feature associated with pts[:,p_idx,:]
|
||||
:param fts: Regional features such that pts[:,p_idx,:] is the feature associated with fts[:,p_idx,:]
|
||||
:return:
|
||||
: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(), self.x_conv.K, self.D).cuda()
|
||||
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))
|
||||
|
||||
if False:
|
||||
# Draw neighborhood points, for debugging.
|
||||
t = 10
|
||||
n = 0
|
||||
test_point = rep_pts[n,t,:].cpu().data.numpy()
|
||||
neighborhood = pts_regional[n,t,:,:].cpu().data.numpy()
|
||||
plt.scatter(pts[n][:,0], pts[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()
|
||||
return fts_p
|
||||
|
||||
class RandPointCNN(nn.Module):
|
||||
""" PointCNN with randomly subsampled representative points. """
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
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(*args, **kwargs)
|
||||
self.pointcnn = PointCNN(C_in, C_out, dims, K, D, P, r_indices_func)
|
||||
self.P = P
|
||||
|
||||
# This is safe because PointCNN requires P.
|
||||
self.P = args[5] if len(args) > 5 else kwargs['P']
|
||||
|
||||
def forward(self, x :Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
|
||||
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.
|
||||
@@ -263,39 +241,3 @@ class RandPointCNN(nn.Module):
|
||||
rep_pts = pts
|
||||
rep_pts_fts = self.pointcnn((rep_pts, pts, fts))
|
||||
return rep_pts, rep_pts_fts
|
||||
|
||||
if __name__ == "__main__":
|
||||
np.random.seed(0)
|
||||
|
||||
N = 1
|
||||
num_points = 1000
|
||||
dims = 2
|
||||
C_in = 4
|
||||
K = 10
|
||||
D = 1
|
||||
|
||||
layer1 = RandPointCNN(C_in, 8, dims, K, D, 1000, knn_indices_func).cuda()
|
||||
layer2 = RandPointCNN( 8, 16, dims, K, D, 500, knn_indices_func).cuda()
|
||||
layer3 = RandPointCNN( 16, 32, dims, K, D, 250, knn_indices_func).cuda()
|
||||
layer4 = RandPointCNN( 32, 64, dims, K, D, 125, knn_indices_func).cuda()
|
||||
layer5 = RandPointCNN( 64, 128, dims, K, D, 50, knn_indices_func).cuda()
|
||||
|
||||
pts = np.random.rand(N,num_points,dims).astype(np.float32)
|
||||
fts = np.random.rand(N,num_points,C_in).astype(np.float32)
|
||||
pts = Variable(torch.from_numpy(pts)).cuda()
|
||||
fts = Variable(torch.from_numpy(fts)).cuda()
|
||||
|
||||
if True:
|
||||
pts, fts = layer1((pts, fts))
|
||||
else:
|
||||
plot(pts, fts)
|
||||
pts, fts = layer1((pts, fts))
|
||||
plot(pts, fts)
|
||||
pts, fts = layer2((pts, fts))
|
||||
plot(pts, fts)
|
||||
pts, fts = layer3((pts, fts))
|
||||
plot(pts, fts)
|
||||
pts, fts = layer4((pts, fts))
|
||||
plot(pts, fts)
|
||||
pts, fts = layer5((pts, fts))
|
||||
plot(pts, fts)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -1,4 +0,0 @@
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from lib import pointcnn
|
||||
+22
-48
@@ -1,3 +1,8 @@
|
||||
"""
|
||||
I got tired of cleaning the code base, so this file will stay as is
|
||||
probably, unless I really want it to be cleaner.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import math
|
||||
@@ -11,27 +16,21 @@ from torch import nn
|
||||
from torch.autograd import Variable
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
import matplotlib.pyplot as plt
|
||||
from mpl_toolkits.mplot3d import Axes3D
|
||||
|
||||
from pointcnn.core import rPointCNN
|
||||
from pointcnn.util import knn_indices_func_gpu
|
||||
from pointcnn.layers import Dense
|
||||
from visualize import *
|
||||
from PointCNN import RandPointCNN
|
||||
from PointCNN import knn_indices_func_gpu
|
||||
from PointCNN.core.util_layers import Dense
|
||||
|
||||
from PointCNN.mnist.visualize import make_dot
|
||||
|
||||
random.seed(0)
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
x = 2
|
||||
|
||||
# N_neighbors, dilution, N_rep, C_out
|
||||
# 8 , 1, all , 16 * x
|
||||
# 8 , 2, all , 32 * x
|
||||
# 8 , 4, all , 48 * x
|
||||
# 12 , 4, 120 , 64 * x
|
||||
# 12 , 6, 120 , 80 * x
|
||||
|
||||
# Data_dim = 3
|
||||
# C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2
|
||||
# (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep)
|
||||
# Abbreviated PointCNN constructor.
|
||||
AbbPointCNN = lambda a,b,c,d,e: RandPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
|
||||
|
||||
class mnist_dataset(Dataset):
|
||||
|
||||
@@ -41,25 +40,21 @@ class mnist_dataset(Dataset):
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
|
||||
def __getitem__(self, i):
|
||||
return self.data[i], self.labels[i]
|
||||
|
||||
# C_in, C_out, D, N_neighbors, dilution, N_rep, r_indices_func, C_lifted = None, mlp_width = 2
|
||||
# (a, b, c, d, e) == (C_in, C_out, N_neighbors, dilution, N_rep)
|
||||
paPointCNN = lambda a,b,c,d,e: rPointCNN(a, b, 3, c, d, e, knn_indices_func_gpu)
|
||||
|
||||
class Classifier(nn.Module):
|
||||
|
||||
|
||||
def __init__(self):
|
||||
super(Classifier, self).__init__()
|
||||
|
||||
self.pcnn1 = paPointCNN( 1, 32, 8, 1, -1)
|
||||
self.pcnn1 = AbbPointCNN( 1, 32, 8, 1, -1)
|
||||
self.pcnn2 = nn.Sequential(
|
||||
paPointCNN( 32, 64, 8, 2, -1),
|
||||
paPointCNN( 64, 96, 8, 4, -1),
|
||||
paPointCNN( 96, 128, 12, 4, 120),
|
||||
paPointCNN(128, 160, 12, 6, 120)
|
||||
AbbPointCNN( 32, 64, 8, 2, -1),
|
||||
AbbPointCNN( 64, 96, 8, 4, -1),
|
||||
AbbPointCNN( 96, 128, 12, 4, 120),
|
||||
AbbPointCNN(128, 160, 12, 6, 120)
|
||||
)
|
||||
|
||||
self.fcn = nn.Sequential(
|
||||
@@ -85,27 +80,6 @@ class Classifier(nn.Module):
|
||||
logits_mean = torch.mean(logits, dim = 1)
|
||||
return logits_mean
|
||||
|
||||
"""
|
||||
def get_indices(batch_size, sample_num, point_num, random_sample = True):
|
||||
indices = []
|
||||
for i in range(batch_size):
|
||||
|
||||
if random_sample:
|
||||
# point_num >= sample_num generally
|
||||
choices = np.random.choice(point_num, sample_num, replace = (point_num < sample_num))
|
||||
else:
|
||||
# This modulo generally not used.
|
||||
choices = np.arange(sample_num) % point_num
|
||||
|
||||
choices = np.expand_dims(choices, axis = 0)
|
||||
b_idx_mat = np.full_like(choices, i)
|
||||
|
||||
# Each set of choices is paired with its batch index.
|
||||
choices_2d = np.concatenate((b_idx_max, choices), axis = 0)
|
||||
indices.append(choices_2d)
|
||||
return np.stack(indices, axis = 1) # (2, batch_size,
|
||||
"""
|
||||
|
||||
model = Classifier().cuda()
|
||||
|
||||
num_class = 10
|
||||
@@ -214,7 +188,7 @@ for e in range(1, num_epochs + 1):
|
||||
loss = loss_fn(out, Variable(label.long()).cuda())
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
|
||||
if global_step % 25 == 0:
|
||||
loss_v = loss.data[0]
|
||||
print("Loss:", loss_v)
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
import h5py
|
||||
|
||||
f = h5py.File("./mnist/zips/train_0.h5", 'r')
|
||||
|
||||
data = f["data"]
|
||||
label = f["label"]
|
||||
|
||||
for r in data[0]:
|
||||
print(r)
|
||||
@@ -1,136 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
pass
|
||||
# from .util import end_channels
|
||||
except:
|
||||
# from util import end_channels
|
||||
pass
|
||||
|
||||
def end_channels(f, make_contiguous = False):
|
||||
class wrapped_layer(nn.Module):
|
||||
def __init__(self):
|
||||
super(wrapped_layer, self).__init__()
|
||||
self.f = f
|
||||
def forward(self, x):
|
||||
x = x.permute(0,3,1,2)
|
||||
x = self.f(x)
|
||||
x = x.permute(0,2,3,1)
|
||||
return x
|
||||
return wrapped_layer()
|
||||
|
||||
class Dense(nn.Module):
|
||||
|
||||
def __init__(self, in_features, out_features, drop_rate = 0, with_bn = True,
|
||||
activation = nn.ReLU()):
|
||||
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):
|
||||
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):
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size, with_bn = True,
|
||||
activation = nn.ReLU()):
|
||||
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):
|
||||
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
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size, depth_multiplier = 1,
|
||||
with_bn = True, activation = nn.ReLU()):
|
||||
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):
|
||||
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, D, *args, **kwargs):
|
||||
super(LayerNorm, self).__init__()
|
||||
if D == 1:
|
||||
self.bn = nn.BatchNorm1d(N, *args, **kwargs)
|
||||
elif D == 2:
|
||||
self.bn = nn.BatchNorm2d(N, *args, **kwargs)
|
||||
elif D == 3:
|
||||
self.bn = nn.BatchNorm3d(N, *args, **kwargs)
|
||||
else:
|
||||
raise ValueError("Dimensionality %i not supported" % D)
|
||||
|
||||
self.forward = lambda x: self.bn(x.unsqueeze(0)).squeeze(0)
|
||||
|
||||
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,
|
||||
LayerNorm(D = 2, 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)
|
||||
@@ -1,118 +0,0 @@
|
||||
import sys, os
|
||||
|
||||
import torch
|
||||
|
||||
import numpy as np
|
||||
from sklearn.neighbors import NearestNeighbors
|
||||
|
||||
torch.CUDA_LAUNCH_BLOCKING = 1
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
sys.path.append(os.path.join(CURRENT_DIR, "..", "lib"))
|
||||
|
||||
import pytorch_knn_cuda
|
||||
|
||||
try:
|
||||
# from .context import pytorch_knn_cuda
|
||||
pass
|
||||
except SystemError:
|
||||
# from context import pytorch_knn_cuda
|
||||
pass
|
||||
|
||||
def apply_along_dim(xs, f, dim):
|
||||
"""
|
||||
PyTorch analog to np.apply_along_axis.
|
||||
:param xs:
|
||||
:param dim:
|
||||
"""
|
||||
return torch.stack([f(x) for x in torch.unbind(xs, dim)], dim)
|
||||
|
||||
def zipwith_matmul(xs, ys):
|
||||
"""
|
||||
Given two lists of 2D matrices of appropriate size, zips them
|
||||
together with matrix multiplication.
|
||||
:param xs:
|
||||
:param ys:
|
||||
"""
|
||||
# xs of shape [N, n, m]
|
||||
# ys of shape [N, m, p]
|
||||
# return shape [N, n, p]
|
||||
N = len(xs)
|
||||
return torch.stack([torch.mm(xs[i], ys[i]) for i in range(N)], dim = 0)
|
||||
|
||||
def knn_indices_func(ps, P, k, d):
|
||||
"""
|
||||
Indexing function based on K-Nearest Neighbors search.
|
||||
:type ps: FloatTensor (N, N_rep, D)
|
||||
:type P: FloatTensor (N, *, D)
|
||||
:type k: int
|
||||
:rtype: FloatTensor (N, N_rep, N_neighbors)
|
||||
:param ps: Representative point
|
||||
:param P: Point cloud to get indices from
|
||||
:param k: Number of nearest neighbors to collect.
|
||||
:return: Array of indices, P_idx, into P such that P[n][P_idx[n],:]
|
||||
is the set k-nearest neighbors for the representative points in P[n].
|
||||
"""
|
||||
ps = ps.data.numpy()
|
||||
P = P.data.numpy()
|
||||
|
||||
def single_batch_knn(p, P_particular):
|
||||
nbrs = NearestNeighbors(d*k + 1, algorithm = "ball_tree").fit(P_particular)
|
||||
indices = nbrs.kneighbors(p)[1]
|
||||
return indices[:,1::d]
|
||||
|
||||
region_idx = np.stack([
|
||||
single_batch_knn(p, P[n]) for n, p in enumerate(ps)
|
||||
], axis = 0)
|
||||
return torch.from_numpy(region_idx)
|
||||
|
||||
# def knn_indices_func_gpu(ps, P, k, d):
|
||||
#
|
||||
# def single_batch_knn(p, P_particular):
|
||||
# nbrs_f = pytorch_knn_cuda.KNearestNeighbor(d*k + 1)
|
||||
# indices = nbrs_f(P_particular, p)[0]
|
||||
# return indices[:,1::d]
|
||||
#
|
||||
# region_idx = torch.stack([
|
||||
# single_batch_knn(p, P[n]) for n, p in enumerate(ps)
|
||||
# ], dim = 0)
|
||||
# return region_idx
|
||||
|
||||
def knn_indices_func_gpu(ps, P, k, d):
|
||||
|
||||
def single_batch_knn(qry, ref):
|
||||
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)
|
||||
return inds[:,1::d]
|
||||
|
||||
region_idx = torch.stack([
|
||||
single_batch_knn(p, P[n]) for n, p in enumerate(ps)
|
||||
], dim = 0)
|
||||
return region_idx
|
||||
|
||||
def plot(pts, fts):
|
||||
num_F = fts.size()[2]
|
||||
pts = pts[0].data.cpu().numpy()
|
||||
plt.scatter(pts[:,0], pts[:,1], s = num_F, c = "k")
|
||||
plt.savefig("./%i.png" % num_F)
|
||||
plt.cla()
|
||||
|
||||
if __name__ == "__main__":
|
||||
from torch.autograd import Variable
|
||||
N_rep = 1000
|
||||
N = 2
|
||||
num_points = 10000
|
||||
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)).cuda()
|
||||
test_ps = Variable(torch.from_numpy(test_ps)).cuda()
|
||||
out = knn_indices_func_gpu(test_ps, test_P, 5, 3)
|
||||
|
||||
print(out)
|
||||
@@ -1,4 +0,0 @@
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
import pointcnn
|
||||
@@ -0,0 +1,82 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from torch.autograd import Variable
|
||||
import numpy as np
|
||||
|
||||
from PointCNN import XConv, RandPointCNN, knn_indices_func_cpu
|
||||
from PointCNN.tests.util_funcs import plot_pts_and_fts
|
||||
|
||||
np.random.seed(0)
|
||||
|
||||
class BasicTests(unittest.TestCase):
|
||||
""" Basic test cases """
|
||||
|
||||
def test_xconv_shape(self):
|
||||
self.assertTrue(True)
|
||||
|
||||
N = 4
|
||||
D = 3
|
||||
C_in = 8
|
||||
C_out = 32
|
||||
N_neighbors = 100
|
||||
|
||||
model = XConv(C_in, C_out, D, N_neighbors).cuda()
|
||||
p = Variable(torch.from_numpy(np.random.rand(N,D).astype(np.float32))).cuda()
|
||||
P = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,D).astype(np.float32))).cuda()
|
||||
F = Variable(torch.from_numpy(np.random.rand(N,N_neighbors,C_in).astype(np.float32))).cuda()
|
||||
out = model(p, P, F)
|
||||
self.assertEqual(out.size(), (N, C_out))
|
||||
|
||||
def test_knn(self):
|
||||
P = np.array([[[0,0],
|
||||
[0,0.95],
|
||||
[1,0],
|
||||
[1,1]]])
|
||||
ps = P[:,[0,3],:]
|
||||
|
||||
P = Variable(torch.from_numpy(P))
|
||||
ps = Variable(torch.from_numpy(ps))
|
||||
|
||||
out = knn_indices_func_cpu(ps, P, 2).numpy()
|
||||
target = np.array([[[1,2],
|
||||
[2,1]]])
|
||||
|
||||
self.assertTrue(np.array_equal(target, out))
|
||||
|
||||
def test_pointcnn_shape(self):
|
||||
N = 1
|
||||
num_points = 1000
|
||||
dims = 2
|
||||
C_in = 4
|
||||
K = 10
|
||||
D = 1
|
||||
|
||||
layer1 = RandPointCNN(C_in, 8, dims, K, D, 1000, knn_indices_func_cpu).cuda()
|
||||
layer2 = RandPointCNN( 8, 16, dims, K, D, 500, knn_indices_func_cpu).cuda()
|
||||
layer3 = RandPointCNN( 16, 32, dims, K, D, 250, knn_indices_func_cpu).cuda()
|
||||
layer4 = RandPointCNN( 32, 64, dims, K, D, 125, knn_indices_func_cpu).cuda()
|
||||
layer5 = RandPointCNN( 64, 128, dims, K, D, 50, knn_indices_func_cpu).cuda()
|
||||
|
||||
pts = np.random.rand(N,num_points,dims).astype(np.float32)
|
||||
fts = np.random.rand(N,num_points,C_in).astype(np.float32)
|
||||
pts = Variable(torch.from_numpy(pts)).cuda()
|
||||
fts = Variable(torch.from_numpy(fts)).cuda()
|
||||
|
||||
if True:
|
||||
pts, fts = layer1((pts, fts))
|
||||
else:
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer1((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer2((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer3((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer4((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer5((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,10 +0,0 @@
|
||||
import unittest
|
||||
|
||||
class AdvancedTests(unittest.TestCase):
|
||||
""" Basic test cases """
|
||||
|
||||
def test_example(self):
|
||||
self.assertTrue(True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+38
-3
@@ -4,14 +4,16 @@ import torch
|
||||
from torch.autograd import Variable
|
||||
import numpy as np
|
||||
|
||||
from pointcnn.core import XConv, knn_indices_func
|
||||
from PointCNN import XConv, RandPointCNN, knn_indices_func_cpu
|
||||
from PointCNN.tests import plot_pts_and_fts
|
||||
|
||||
np.random.seed(0)
|
||||
|
||||
class BasicTests(unittest.TestCase):
|
||||
""" Basic test cases """
|
||||
|
||||
def test_xconv_shape(self):
|
||||
self.assertTrue(True)
|
||||
np.random.seed(0)
|
||||
|
||||
N = 4
|
||||
D = 3
|
||||
@@ -36,12 +38,45 @@ class BasicTests(unittest.TestCase):
|
||||
P = Variable(torch.from_numpy(P))
|
||||
ps = Variable(torch.from_numpy(ps))
|
||||
|
||||
out = knn_indices_func(ps, P, 2).numpy()
|
||||
out = knn_indices_func_cpu(ps, P, 2).numpy()
|
||||
target = np.array([[[1,2],
|
||||
[2,1]]])
|
||||
|
||||
self.assertTrue(np.array_equal(target, out))
|
||||
|
||||
def test_pointcnn_shape(self):
|
||||
N = 1
|
||||
num_points = 1000
|
||||
dims = 2
|
||||
C_in = 4
|
||||
K = 10
|
||||
D = 1
|
||||
|
||||
layer1 = RandPointCNN(C_in, 8, dims, K, D, 1000, knn_indices_func_cpu).cuda()
|
||||
layer2 = RandPointCNN( 8, 16, dims, K, D, 500, knn_indices_func_cpu).cuda()
|
||||
layer3 = RandPointCNN( 16, 32, dims, K, D, 250, knn_indices_func_cpu).cuda()
|
||||
layer4 = RandPointCNN( 32, 64, dims, K, D, 125, knn_indices_func_cpu).cuda()
|
||||
layer5 = RandPointCNN( 64, 128, dims, K, D, 50, knn_indices_func_cpu).cuda()
|
||||
|
||||
pts = np.random.rand(N,num_points,dims).astype(np.float32)
|
||||
fts = np.random.rand(N,num_points,C_in).astype(np.float32)
|
||||
pts = Variable(torch.from_numpy(pts)).cuda()
|
||||
fts = Variable(torch.from_numpy(fts)).cuda()
|
||||
|
||||
if True:
|
||||
pts, fts = layer1((pts, fts))
|
||||
else:
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer1((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer2((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer3((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer4((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
pts, fts = layer5((pts, fts))
|
||||
plot_pts_and_fts(pts, fts)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
# External Modules
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
# Internal Modules
|
||||
from PointCNN.core import UFloatTensor
|
||||
|
||||
def plot_pts_and_fts(pts : UFloatTensor, # (N, x, dims)
|
||||
fts : UFloatTensor # (N, x, y)
|
||||
) -> None:
|
||||
"""
|
||||
Visualization function. Shows points and number of features, represented by
|
||||
the size of the point.
|
||||
:param pts: Point cloud such that fts[:,p_idx,:] is the feature associated
|
||||
with pts[:,p_idx,:].
|
||||
:param fts: Features such that pts[:,p_idx,:] is the feature associated
|
||||
with fts[:,p_idx,:].
|
||||
"""
|
||||
if pts.is_cuda:
|
||||
pts = pts.cpu()
|
||||
num_F = fts.size()[2]
|
||||
pts = pts[0].data.numpy()
|
||||
plt.scatter(pts[:,0], pts[:,1], s = num_F, c = "k")
|
||||
plt.show()
|
||||
plt.cla()
|
||||
|
||||
def plot_neighborhood(pts : UFloatTensor, # (N, x, dims)
|
||||
rep_pts : UFloatTensor, # (N, P, dims)
|
||||
pts_regional : UFloatTensor # (N, P, dims)
|
||||
) -> None:
|
||||
"""
|
||||
Visualization function. Shows neighborhood points around a randomly
|
||||
selected representative.
|
||||
:param pts: Point cloud.
|
||||
:param rep_pts: Representative points.
|
||||
:param pts_regional: Regional neighborhoods around representative points.
|
||||
"""
|
||||
if rep_pts.is_cuda:
|
||||
rep_pts = rep_pts.cpu()
|
||||
pts_regional = pts_regional.cpu()
|
||||
n = np.randint(0, rep_pts.shape[0])
|
||||
t = np.randint(0, rep_pts.shape[1])
|
||||
test_point = rep_pts[n,t,:].data.numpy()
|
||||
neighborhood = pts_regional[n,t,:,:].data.numpy()
|
||||
plt.scatter(pts[n][:,0], pts[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()
|
||||
plt.cla()
|
||||
Reference in New Issue
Block a user