wassname/Pointnet2_PyTorch
1
0
Fork 0
mirror of https://github.com/wassname/Pointnet2_PyTorch.git synced 2026-06-27 16:00:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Updates

Browse Source
This commit is contained in:
erikwijmans
2018-01-06 12:13:52 -05:00
parent 7e746ba72a
commit 5a5adc2b77
20 changed files with 650 additions and 494 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/build_ffi.py
+7 -3
View File
@@ -9,7 +9,8 @@ base_dir = osp.dirname(osp.abspath(__file__))
def parse_args():
parser = argparse.ArgumentParser(
description="Arguments for building pointnet2 ffi extension")
description="Arguments for building pointnet2 ffi extension"
)
parser.add_argument("--objs", nargs="*")
clean_arg = parser.add_mutually_exclusive_group()
clean_arg.add_argument("--build", dest='build', action="store_true")
@@ -27,7 +28,7 @@ def build(args):
extra_objects += [a for a in glob.glob('/usr/local/cuda/lib64/*.a')]
ffi = create_extension(
'_ext',
'_ext.pointnet2',
headers=[a for a in glob.glob("cinclude/*_wrapper.h")],
sources=[a for a in glob.glob("csrc/*.c")],
define_macros=[('WITH_CUDA', None)],
@@ -36,12 +37,15 @@ def build(args):
extra_objects=extra_objects,
include_dirs=[osp.join(base_dir, 'cinclude')],
verbose=False,
package=False)
package=False
)
ffi.build()
def clean(args):
shutil.rmtree(osp.join(base_dir, "_ext"))
if __name__ == "__main__":
args = parse_args()
if args.clean:
utils/csrc/ball_query.c
+7 -7
View File
@@ -8,13 +8,13 @@ int ball_query_wrapper(int b, int n, int m, float radius, int nsample,
THCudaTensor *new_xyz_tensor, THCudaTensor *xyz_tensor,
THCudaIntTensor *idx_tensor) {
const float *new_xyz = THCudaTensor_data(state, new_xyz_tensor);
const float *xyz = THCudaTensor_data(state, xyz_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *new_xyz = THCudaTensor_data(state, new_xyz_tensor);
const float *xyz = THCudaTensor_data(state, xyz_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
query_ball_point_kernel_wrapper(b, n, m, radius, nsample, new_xyz, xyz,
idx, stream);
return 1;
query_ball_point_kernel_wrapper(b, n, m, radius, nsample, new_xyz, xyz, idx,
stream);
return 1;
}
utils/csrc/ball_query_gpu.cu
+35 -37
View File
@@ -11,38 +11,37 @@ __global__ void query_ball_point_kernel(int b, int n, int m, float radius,
int nsample,
const float *__restrict__ new_xyz,
const float *__restrict__ xyz,
int * __restrict__ idx) {
int batch_index = blockIdx.x;
xyz += batch_index * n * 3;
new_xyz += batch_index * m * 3;
idx += m * nsample * batch_index;
int *__restrict__ idx) {
int batch_index = blockIdx.x;
xyz += batch_index * n * 3;
new_xyz += batch_index * m * 3;
idx += m * nsample * batch_index;
int index = threadIdx.x;
int stride = blockDim.x;
int index = threadIdx.x;
int stride = blockDim.x;
float radius2 = radius * radius;
for (int j = index; j < m; j += stride) {
float new_x = new_xyz[j * 3 + 0];
float new_y = new_xyz[j * 3 + 1];
float new_z = new_xyz[j * 3 + 2];
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
float x = xyz[k * 3 + 0];
float y = xyz[k * 3 + 1];
float z = xyz[k * 3 + 2];
float d2 = (new_x - x) * (new_x - x) +
(new_y - y) * (new_y - y) +
(new_z - z) * (new_z - z);
if (d2 < radius2) {
if (cnt == 0) {
for (int l = 0; l < nsample; ++l) {
idx[j * nsample + l] = k;
}
}
idx[j * nsample + cnt] = k;
++cnt;
}
float radius2 = radius * radius;
for (int j = index; j < m; j += stride) {
float new_x = new_xyz[j * 3 + 0];
float new_y = new_xyz[j * 3 + 1];
float new_z = new_xyz[j * 3 + 2];
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
float x = xyz[k * 3 + 0];
float y = xyz[k * 3 + 1];
float z = xyz[k * 3 + 2];
float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) +
(new_z - z) * (new_z - z);
if (d2 < radius2) {
if (cnt == 0) {
for (int l = 0; l < nsample; ++l) {
idx[j * nsample + l] = k;
}
}
idx[j * nsample + cnt] = k;
++cnt;
}
}
}
}
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
@@ -50,14 +49,13 @@ void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
const float *xyz, int *idx,
cudaStream_t stream) {
cudaError_t err;
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
b, n, m, radius, nsample, new_xyz, xyz, idx);
cudaError_t err;
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
b, n, m, radius, nsample, new_xyz, xyz, idx);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n",
cudaGetErrorString(err));
exit(-1);
}
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
}
utils/csrc/group_points.c
+14 -14
View File
@@ -9,15 +9,15 @@ int group_points_wrapper(int b, int n, int c, int npoints, int nsample,
THCudaIntTensor *idx_tensor,
THCudaTensor *out_tensor) {
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
group_points_kernel_wrapper(b, n, c, npoints, nsample, points, idx, out,
stream);
return 1;
group_points_kernel_wrapper(b, n, c, npoints, nsample, points, idx, out,
stream);
return 1;
}
int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
@@ -25,13 +25,13 @@ int group_points_grad_wrapper(int b, int n, int c, int npoints, int nsample,
THCudaIntTensor *idx_tensor,
THCudaTensor *grad_points_tensor) {
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
group_points_grad_kernel_wrapper(b, n, c, npoints, nsample, grad_out,
idx, grad_points, stream);
return 1;
group_points_grad_kernel_wrapper(b, n, c, npoints, nsample, grad_out, idx,
grad_points, stream);
return 1;
}
utils/csrc/group_points_gpu.cu
+49 -52
View File
@@ -1,86 +1,83 @@
#include <stdio.h>
#include <stdlib.h>
#include "group_points_gpu.h"
#include "cuda_utils.h"
#include "group_points_gpu.h"
// input: points(b, n, c) idx(b, npoints, nsample)
// output: out(b, npoints, nsample, c)
__global__ void group_points_kernel(int b, int n, int c, int npoints,
int nsample,
const float *__restrict__ points,
const int *__restrict__ idx,
float *__restrict__ out) {
int batch_index = blockIdx.x;
points += batch_index * n * c;
idx += batch_index * npoints * nsample;
out += batch_index * npoints * nsample * c;
const float *__restrict__ points,
const int *__restrict__ idx,
float *__restrict__ out) {
int batch_index = blockIdx.x;
points += batch_index * n * c;
idx += batch_index * npoints * nsample;
out += batch_index * npoints * nsample * c;
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
memcpy(out + j * nsample * c + k * c, points + ii * c,
sizeof(float) * c);
}
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
memcpy(out + j * nsample * c + k * c, points + ii * c,
sizeof(float) * c);
}
}
}
void group_points_kernel_wrapper(int b, int n, int c, int npoints, int nsample,
const float *points, const int *idx,
float *out, cudaStream_t stream) {
cudaError_t err;
group_points_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, points, idx, out);
cudaError_t err;
group_points_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, points, idx, out);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n",
cudaGetErrorString(err));
exit(-1);
}
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
}
// input: grad_out(b, npoints, nsample, c), idx(b, npoints, nsample)
// output: grad_points(b, n, c)
__global__ void group_points_grad_kernel(int b, int n, int c, int npoints,
int nsample,
const float *__restrict__ grad_out,
const int *__restrict__ idx,
float *__restrict__ grad_points) {
int batch_index = blockIdx.x;
grad_points += batch_index * n * c;
idx += batch_index * npoints * nsample;
grad_out += batch_index * npoints * nsample * c;
const float *__restrict__ grad_out,
const int *__restrict__ idx,
float *__restrict__ grad_points) {
int batch_index = blockIdx.x;
grad_points += batch_index * n * c;
idx += batch_index * npoints * nsample;
grad_out += batch_index * npoints * nsample * c;
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
for (int l = 0; l < c; ++l) {
atomicAdd(
grad_points + ii * c + l,
grad_out[j * nsample * c + k * c + l]);
}
}
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < npoints; j += stride) {
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
for (int l = 0; l < c; ++l) {
atomicAdd(grad_points + ii * c + l,
grad_out[j * nsample * c + k * c + l]);
}
}
}
}
void group_points_grad_kernel_wrapper(int b, int n, int c, int npoints,
int nsample, const float *grad_out,
const int *idx, float *grad_points,
cudaStream_t stream) {
cudaError_t err;
group_points_grad_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, grad_out, idx, grad_points);
cudaError_t err;
group_points_grad_kernel<<<b, opt_n_threads(npoints), 0, stream>>>(
b, n, c, npoints, nsample, grad_out, idx, grad_points);
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n",
cudaGetErrorString(err));
exit(-1);
}
err = cudaGetLastError();
if (cudaSuccess != err) {
fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
exit(-1);
}
}
utils/csrc/interpolate_gpu.cu
+1 -1
View File
@@ -2,8 +2,8 @@
#include <stdio.h>
#include <stdlib.h>
#include "interpolate_gpu.h"
#include "cuda_utils.h"
#include "interpolate_gpu.h"
// input: unknown(b, n, 3) known(b, m, 3)
// output: dist2(b, n, 3), idx(b, n, 3)
utils/csrc/roi_mask_points_wrapper.c
+28 -29
View File
@@ -9,17 +9,17 @@ int roi_mask_wrapper(int n_roi, int b, int n, THCudaTensor *rois_tensor,
THCudaTensor *data_xyz_tensor,
THCudaByteTensor *mask_tensor) {
const float *rois = THCudaTensor_data(state, rois_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const float *data_xyz = THCudaTensor_data(state, data_xyz_tensor);
unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *rois = THCudaTensor_data(state, rois_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const float *data_xyz = THCudaTensor_data(state, data_xyz_tensor);
unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_mask_kernel_wrapper(n_roi, b, n, rois, batch_indices, data_xyz,
mask, stream);
return 1;
roi_mask_kernel_wrapper(n_roi, b, n, rois, batch_indices, data_xyz, mask,
stream);
return 1;
}
int roi_avg_pool_forward_wrapper(int n_roi, int b, int n, int d,
@@ -28,17 +28,17 @@ int roi_avg_pool_forward_wrapper(int n_roi, int b, int n, int d,
THCudaTensor *points_tensor,
THCudaTensor *descriptors_tensor) {
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
float *descriptors = THCudaTensor_data(state, descriptors_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
float *descriptors = THCudaTensor_data(state, descriptors_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_forward_wrapper(n_roi, b, n, d, mask, batch_indices,
points, descriptors, stream);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_forward_wrapper(n_roi, b, n, d, mask, batch_indices,
points, descriptors, stream);
return 1;
return 1;
}
int roi_avg_pool_backward_wrapper(int n_roi, int b, int n, int d,
@@ -47,17 +47,16 @@ int roi_avg_pool_backward_wrapper(int n_roi, int b, int n, int d,
THCudaTensor *grad_descriptors_tensor,
THCudaTensor *grad_points_tensor) {
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *grad_descriptors =
THCudaTensor_data(state, grad_descriptors_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
const long *batch_indices =
THCudaLongTensor_data(state, batch_indices_tensor);
const unsigned char *mask = THCudaByteTensor_data(state, mask_tensor);
const float *grad_descriptors =
THCudaTensor_data(state, grad_descriptors_tensor);
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_backward_wrapper(n_roi, b, n, d, mask,
batch_indices, grad_descriptors,
grad_points, stream);
cudaStream_t stream = THCState_getCurrentStream(state);
roi_avg_pool_kernel_backward_wrapper(n_roi, b, n, d, mask, batch_indices,
grad_descriptors, grad_points, stream);
return 1;
return 1;
}
utils/csrc/sampling.c
+12 -14
View File
@@ -9,15 +9,14 @@ int gather_points_wrapper(int b, int n, int c, int npoints,
THCudaIntTensor *idx_tensor,
THCudaTensor *out_tensor) {
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
const int *idx = THCudaIntTensor_data(state, idx_tensor);
float *out = THCudaTensor_data(state, out_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out,
stream);
return 1;
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out, stream);
return 1;
}
int furthest_point_sampling_wrapper(int b, int n, int m,
@@ -25,13 +24,12 @@ int furthest_point_sampling_wrapper(int b, int n, int m,
THCudaTensor *temp_tensor,
THCudaIntTensor *idx_tensor) {
const float *points = THCudaTensor_data(state, points_tensor);
float *temp = THCudaTensor_data(state, temp_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
const float *points = THCudaTensor_data(state, points_tensor);
float *temp = THCudaTensor_data(state, temp_tensor);
int *idx = THCudaIntTensor_data(state, idx_tensor);
cudaStream_t stream = THCState_getCurrentStream(state);
cudaStream_t stream = THCState_getCurrentStream(state);
furthest_point_sampling_kernel_wrapper(b, n, m, points, temp, idx,
stream);
return 1;
furthest_point_sampling_kernel_wrapper(b, n, m, points, temp, idx, stream);
return 1;
}
utils/data_utils.py
+11 -4
View File
@@ -3,17 +3,19 @@ import numpy as np
class PointcloudScale(object):
def __init__(self, mean=2.0, std=1.0, clip=1.8):
self.mean, self.std, self.clip = mean, std, clip
def __call__(self, points):
scaler = points.new(1).normal_(
mean=self.mean, std=self.std).clamp_(
max(self.mean - self.clip, 0.01), self.mean + self.clip)
mean=self.mean, std=self.std
).clamp_(max(self.mean - self.clip, 0.01), self.mean + self.clip)
return scaler * points
class PointcloudRotate(object):
def __init__(self, x_axis=False, z_axis=True):
assert x_axis or z_axis
self.x, self.z = x_axis, z_axis
@@ -46,25 +48,30 @@ class PointcloudRotate(object):
class PointcloudJitter(object):
def __init__(self, std=0.01, clip=0.03):
self.std, self.clip = std, clip
def __call__(self, points):
jittered_data = points.new(*points.size()).normal_(
mean=0.0, std=self.std).clamp_(-self.clip, self.clip)
mean=0.0, std=self.std
).clamp_(-self.clip, self.clip)
return points + jittered_data
class PointcloudTranslate(object):
def __init__(self, std=1.0, clip=3.0):
self.std, self.clip = std, clip
def __call__(self, points):
translation = points.new(3).normal_(
mean=0.0, std=self.std).clamp_(-self.clip, self.clip)
mean=0.0, std=self.std
).clamp_(-self.clip, self.clip)
return points + translation
class PointcloudToTensor(object):
def __call__(self, points):
return torch.from_numpy(points).float()
utils/linalg_utils.py
+5 -3
View File
@@ -4,9 +4,11 @@ from enum import Enum
PDist2Order = Enum('PDist2Order', 'd_first d_second')
def pdist2(X: torch.Tensor,
Z: torch.Tensor = None,
order: PDist2Order = PDist2Order.d_second) -> torch.Tensor:
def pdist2(
X: torch.Tensor,
Z: torch.Tensor = None,
order: PDist2Order = PDist2Order.d_second
) -> torch.Tensor:
r""" Calculates the pairwise distance between X and Z
D[b, i, j] = l2 distance X[b, i] and Z[b, j]
utils/pointnet2_modules.py
+54 -44
View File
@@ -24,13 +24,15 @@ class PointnetSAModuleMSG(nn.Module):
Use batchnorm
"""
def __init__(self,
*,
npoint: int,
radii: List[float],
nsamples: List[int],
mlps: List[List[int]],
bn: bool = True):
def __init__(
self,
*,
npoint: int,
radii: List[float],
nsamples: List[int],
mlps: List[List[int]],
bn: bool = True
):
super().__init__()
assert len(radii) == len(nsamples) == len(mlps)
@@ -41,8 +43,7 @@ class PointnetSAModuleMSG(nn.Module):
for i in range(len(radii)):
radius = radii[i]
nsample = nsamples[i]
self.groupers.append(
pointnet2_utils.QueryAndGroup(radius, nsample))
self.groupers.append(pointnet2_utils.QueryAndGroup(radius, nsample))
mlp_spec = mlps[i]
self.mlps.append(pt_utils.SharedMLP(mlp_spec, bn=bn))
@@ -66,18 +67,20 @@ class PointnetSAModuleMSG(nn.Module):
new_points_list = []
new_xyz = pointnet2_utils.gather_points(
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint))
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
)
for i in range(len(self.groupers)):
new_points = self.groupers[i](xyz, new_xyz, points)
new_points = self.mlps[i](new_points.permute(
0, 3, 1, 2)) # (B, mlp[-1], npoint, nsample)
new_points = self.mlps[i](new_points.permute(0, 3, 1, 2)
) # (B, mlp[-1], npoint, nsample)
new_points = F.max_pool2d(
new_points,
kernel_size=[1, new_points.size(3)]) # (B, mlp[-1], npoint, 1)
new_points, kernel_size=[1, new_points.size(3)]
) # (B, mlp[-1], npoint, 1)
new_points = new_points.squeeze(-1) # (B, mlp[-1], npoint)
new_points = new_points.transpose(
1, 2).contiguous() # (B, npoint, mlp[-1])
1, 2
).contiguous() # (B, npoint, mlp[-1])
new_points_list.append(new_points)
@@ -101,13 +104,15 @@ class PointnetSAModule(nn.Module):
Use batchnorm
"""
def __init__(self,
*,
mlp: List[int],
npoint: int = None,
radius: float = None,
nsample: int = None,
bn: bool = True):
def __init__(
self,
*,
mlp: List[int],
npoint: int = None,
radius: float = None,
nsample: int = None,
bn: bool = True
):
super().__init__()
self.npoint = npoint
@@ -140,21 +145,23 @@ class PointnetSAModule(nn.Module):
if self.npoint is not None:
new_xyz = pointnet2_utils.gather_points(
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint))
xyz, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
)
else:
new_xyz = xyz.data.new([[[0, 0, 0]]]).expand(xyz.size(0), 1, 3)
new_points = self.grouper(xyz, new_xyz,
points) # (B, npoint, nsample, 3 + C)
new_points = self.grouper(
xyz, new_xyz, points
) # (B, npoint, nsample, 3 + C)
new_points = self.mlp(new_points.permute(
0, 3, 1, 2)) # (B, mlp[-1], npoint, nsample)
new_points = self.mlp(new_points.permute(0, 3, 1, 2)
) # (B, mlp[-1], npoint, nsample)
new_points = F.max_pool2d(
new_points,
kernel_size=[1, new_points.size(3)]) # (B, mlp[-1], npoint, 1)
new_points, kernel_size=[1, new_points.size(3)]
) # (B, mlp[-1], npoint, 1)
new_points = new_points.squeeze(-1) # (B, mlp[-1], npoint)
new_points = new_points.transpose(
1, 2).contiguous() # (B, npoint, mlp[-1])
new_points = new_points.transpose(1, 2
).contiguous() # (B, npoint, mlp[-1])
return new_xyz, new_points
@@ -174,9 +181,10 @@ class PointnetFPModule(nn.Module):
super().__init__()
self.mlp = pt_utils.SharedMLP(mlp, bn=bn)
def forward(self, unknown: torch.Tensor, known: torch.Tensor,
unknow_feats: torch.Tensor,
known_feats: torch.Tensor) -> torch.Tensor:
def forward(
self, unknown: torch.Tensor, known: torch.Tensor,
unknow_feats: torch.Tensor, known_feats: torch.Tensor
) -> torch.Tensor:
r"""
Parameters
----------
@@ -201,19 +209,21 @@ class PointnetFPModule(nn.Module):
weight = dist_recip / norm
interpolated_feats = pointnet2_utils.three_interpolate(
known_feats, idx, weight)
known_feats, idx, weight
)
if unknow_feats is not None:
new_points = torch.cat(
[interpolated_feats, unknow_feats], dim=-1) #(B, n, C2 + C1)
new_points = torch.cat([interpolated_feats, unknow_feats],
dim=-1) #(B, n, C2 + C1)
else:
new_points = interpolated_feats
new_points = new_points.unsqueeze(-1).transpose(1,
2) #(B, C2 + C1, n, 1)
new_points = new_points.unsqueeze(-1).transpose(
1, 2
) #(B, C2 + C1, n, 1)
new_points = self.mlp(new_points)
return new_points.squeeze(-1).transpose(
1, 2).contiguous() #(B, n, mlp[-1])
return new_points.squeeze(-1).transpose(1, 2
).contiguous() #(B, n, mlp[-1])
if __name__ == "__main__":
@@ -224,7 +234,8 @@ if __name__ == "__main__":
xyz_feats = Variable(torch.randn(2, 10, 6).cuda(), requires_grad=True)
test_module = PointnetSAModuleMSG(
npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]])
npoint=2, radii=[5.0, 10.0], nsamples=[6, 3], mlps=[[9, 3], [9, 6]]
)
test_module.cuda()
print(test_module(xyz, xyz_feats))
@@ -237,7 +248,6 @@ if __name__ == "__main__":
for _ in range(1):
_, new_points = test_module(xyz, xyz_feats)
new_points.backward(
torch.cuda.FloatTensor(*new_points.size()).fill_(1))
new_points.backward(torch.cuda.FloatTensor(*new_points.size()).fill_(1))
print(new_points)
print(xyz.grad)
utils/pointnet2_utils.py
+57 -36
View File
@@ -5,12 +5,14 @@ import torch.nn.functional as F
import torch.nn as nn
from linalg_utils import pdist2, PDist2Order
from collections import namedtuple
import _ext as pointnet2
import pytorch_utils as pt_utils
from typing import List, Tuple
from _ext import pointnet2
class RandomDropout(nn.Module):
def __init__(self, p=0.5, inplace=False):
super().__init__()
self.p = p
@@ -18,11 +20,13 @@ class RandomDropout(nn.Module):
def forward(self, X):
theta = torch.Tensor(1).uniform_(0, self.p)[0]
return pt_utils.feature_dropout_no_scaling(X, theta, self.train,
self.inplace)
return pt_utils.feature_dropout_no_scaling(
X, theta, self.train, self.inplace
)
class FurthestPointSampling(Function):
@staticmethod
def forward(ctx, xyz: torch.Tensor, npoint: int) -> torch.Tensor:
r"""
@@ -30,16 +34,16 @@ class FurthestPointSampling(Function):
minimum distance
Parameters
---------
----------
xyz : torch.Tensor
(B, N, 3) tensor where N > npoint
npoint : int32
number of points in the sampled set
Returns
-------
torch.Tensor
(B, npoint) tensor containing the set
------
"""
B, N, _ = xyz.size()
@@ -50,8 +54,9 @@ class FurthestPointSampling(Function):
temp = temp.contiguous()
output = output.contiguous()
pointnet2.furthest_point_sampling_wrapper(B, N, npoint, xyz, temp,
output)
pointnet2.furthest_point_sampling_wrapper(
B, N, npoint, xyz, temp, output
)
return output
@@ -64,6 +69,7 @@ furthest_point_sample = FurthestPointSampling.apply
class GatherPoints(Function):
@staticmethod
def forward(ctx, points: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
r"""
@@ -71,7 +77,7 @@ class GatherPoints(Function):
minimum distance
Parameters
---------
----------
points : torch.Tensor
(B, N, 3) tensor
@@ -79,9 +85,9 @@ class GatherPoints(Function):
(B, npoint) tensor of the points to gather
Returns
-------
torch.Tensor
(B, npoint, 3) tensor
------
"""
B, N, C = points.size()
@@ -106,6 +112,7 @@ gather_points = GatherPoints.apply
class ThreeNN(Function):
@staticmethod
def forward(ctx, unknown: torch.Tensor,
known: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -147,9 +154,11 @@ three_nn = ThreeNN.apply
class ThreeInterpolate(Function):
@staticmethod
def forward(ctx, points: torch.Tensor, idx: torch.Tensor,
weight: torch.Tensor) -> torch.Tensor:
def forward(
ctx, points: torch.Tensor, idx: torch.Tensor, weight: torch.Tensor
) -> torch.Tensor:
r"""
Performs weight linear interpolation on 3 points
Parameters
@@ -178,14 +187,15 @@ class ThreeInterpolate(Function):
idx = idx.contiguous()
weight = weight.contiguous()
output = output.contiguous()
pointnet2.three_interpolate_wrapper(B, m, c, n, points, idx, weight,
output)
pointnet2.three_interpolate_wrapper(
B, m, c, n, points, idx, weight, output
)
return output
@staticmethod
def backward(ctx, grad_out: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
r"""
Parameters
----------
@@ -196,6 +206,7 @@ class ThreeInterpolate(Function):
-------
grad_points : torch.Tensor
(B, m, c) tensor with gradients of points
None
None
@@ -209,8 +220,9 @@ class ThreeInterpolate(Function):
idx = idx.contiguous()
weight = weight.contiguous()
grad_points = grad_points.contiguous()
pointnet2.three_interpolate_grad_wrapper(B, n, c, m, grad_out.data,
idx, weight, grad_points.data)
pointnet2.three_interpolate_grad_wrapper(
B, n, c, m, grad_out.data, idx, weight, grad_points.data
)
return grad_points, None, None
@@ -219,6 +231,7 @@ three_interpolate = ThreeInterpolate.apply
class GroupPoints(Function):
@staticmethod
def forward(ctx, points: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
r"""
@@ -243,8 +256,9 @@ class GroupPoints(Function):
points = points.contiguous()
idx = idx.contiguous()
output = output.contiguous()
pointnet2.group_points_wrapper(B, N, C, npoints, nsample, points, idx,
output)
pointnet2.group_points_wrapper(
B, N, C, npoints, nsample, points, idx, output
)
ctx.idx_N_C_for_backward = (idx, N, C)
return output
@@ -273,7 +287,8 @@ class GroupPoints(Function):
grad_out = grad_out.contiguous()
grad_points = grad_points.contiguous()
pointnet2.group_points_grad_wrapper(
B, N, C, npoint, nsample, grad_out.data, idx, grad_points.data)
B, N, C, npoint, nsample, grad_out.data, idx, grad_points.data
)
return grad_points, None
@@ -282,13 +297,16 @@ group_points = GroupPoints.apply
class BallQuery(Function):
@staticmethod
def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor,
new_xyz: torch.Tensor) -> torch.Tensor:
def forward(
ctx, radius: float, nsample: int, xyz: torch.Tensor,
new_xyz: torch.Tensor
) -> torch.Tensor:
r"""
Parameters
---------
----------
radius : float
radius of the balls
nsample : int
@@ -299,7 +317,7 @@ class BallQuery(Function):
(B, npoint, 3) centers of the ball query
Returns
------
-------
torch.Tensor
(B, npoint, nsample) tensor with the indicies of the points that form the query balls
"""
@@ -311,8 +329,9 @@ class BallQuery(Function):
new_xyz = new_xyz.contiguous()
xyz = xyz.contiguous()
idx = idx.contiguous()
pointnet2.ball_query_wrapper(B, N, npoint, radius, nsample, new_xyz,
xyz, idx)
pointnet2.ball_query_wrapper(
B, N, npoint, radius, nsample, new_xyz, xyz, idx
)
return idx
@@ -344,10 +363,11 @@ class QueryAndGroup(nn.Module):
self,
xyz: torch.Tensor,
new_xyz: torch.Tensor,
points: torch.Tensor = None) -> Tuple[torch.Tensor]:
points: torch.Tensor = None
) -> Tuple[torch.Tensor]:
r"""
Parameters
---------
----------
xyz : torch.Tensor
xyz coordinates of the points (B, N, 3)
new_xyz : torch.Tensor
@@ -368,9 +388,8 @@ class QueryAndGroup(nn.Module):
if points is not None:
grouped_points = group_points(points, idx)
if self.use_xyz:
new_points = torch.cat(
[grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
new_points = torch.cat([grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
else:
new_points = group_points
else:
@@ -395,10 +414,11 @@ class GroupAll(nn.Module):
self,
xyz: torch.Tensor,
new_xyz: torch.Tensor,
points: torch.Tensor = None) -> Tuple[torch.Tensor]:
points: torch.Tensor = None
) -> Tuple[torch.Tensor]:
r"""
Parameters
---------
----------
xyz : torch.Tensor
xyz coordinates of the points (B, N, 3)
new_xyz : torch.Tensor
@@ -414,11 +434,12 @@ class GroupAll(nn.Module):
grouped_xyz = xyz.view(xyz.size(0), 1, xyz.size(1), xyz.size(2))
if points is not None:
grouped_points = points.view(points.size(0), 1, points.size(1), points.size(2))
grouped_points = points.view(
points.size(0), 1, points.size(1), points.size(2)
)
if self.use_xyz:
new_points = torch.cat(
[grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
new_points = torch.cat([grouped_xyz, grouped_points],
dim=-1) # (B, npoint, nsample, 3 + C)
else:
new_points = group_points
else:
utils/pytorch_utils.py
+182 -141
View File
@@ -16,48 +16,55 @@ import math
class SharedMLP(nn.Sequential):
def __init__(self,
args: List[int],
*,
bn: bool = False,
activation=nn.ReLU(inplace=True),
name: str = ""):
def __init__(
self,
args: List[int],
*,
bn: bool = False,
activation=nn.ReLU(inplace=True),
name: str = ""
):
super().__init__()
for i in range(len(args) - 1):
self.add_module(name + 'layer{}'.format(i),
Conv2d(
args[i],
args[i + 1],
bn=bn,
activation=activation))
self.add_module(
name + 'layer{}'.format(i),
Conv2d(args[i], args[i + 1], bn=bn, activation=activation)
)
class _ConvBase(nn.Sequential):
def __init__(self,
in_size,
out_size,
kernel_size,
stride,
padding,
activation,
bn,
init,
conv=None,
batch_norm=None,
bias=True,
name=""):
def __init__(
self,
in_size,
out_size,
kernel_size,
stride,
padding,
activation,
bn,
init,
conv=None,
batch_norm=None,
bias=True,
name=""
):
super().__init__()
bias = bias and (not bn)
self.add_module(name + 'conv',
conv(
in_size,
out_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias))
self.add_module(
name + 'conv',
conv(
in_size,
out_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias
)
)
init(self[0].weight)
if bias:
@@ -73,18 +80,21 @@ class _ConvBase(nn.Sequential):
class Conv1d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -97,22 +107,26 @@ class Conv1d(_ConvBase):
conv=nn.Conv1d,
batch_norm=nn.BatchNorm1d,
bias=bias,
name=name)
name=name
)
class Conv2d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int] = (1, 1),
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int] = (1, 1),
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -125,22 +139,26 @@ class Conv2d(_ConvBase):
conv=nn.Conv2d,
batch_norm=nn.BatchNorm2d,
bias=bias,
name=name)
name=name
)
class Conv3d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int, int] = (1, 1, 1),
stride: Tuple[int, int, int] = (1, 1, 1),
padding: Tuple[int, int, int] = (0, 0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int, int] = (1, 1, 1),
stride: Tuple[int, int, int] = (1, 1, 1),
padding: Tuple[int, int, int] = (0, 0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -153,18 +171,22 @@ class Conv3d(_ConvBase):
conv=nn.Conv3d,
batch_norm=nn.BatchNorm3d,
bias=bias,
name=name)
name=name
)
class FC(nn.Sequential):
def __init__(self,
in_size: int,
out_size: int,
*,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=None,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=None,
name: str = ""
):
super().__init__()
self.add_module(name + 'fc', nn.Linear(in_size, out_size, bias=not bn))
if init is not None:
@@ -183,6 +205,7 @@ class FC(nn.Sequential):
class _DropoutNoScaling(InplaceFunction):
@staticmethod
def _make_noise(input):
return input.new().resize_as_(input)
@@ -192,8 +215,9 @@ class _DropoutNoScaling(InplaceFunction):
if inplace:
return None
n = g.appendNode(
g.create("Dropout", [input]).f_("ratio", p).i_(
"is_test", not train))
g.create("Dropout", [input]).f_("ratio",
p).i_("is_test", not train)
)
real = g.appendNode(g.createSelect(n, 0))
g.appendNode(g.createSelect(n, 1))
return real
@@ -201,8 +225,10 @@ class _DropoutNoScaling(InplaceFunction):
@classmethod
def forward(cls, ctx, input, p=0.5, train=False, inplace=False):
if p < 0 or p > 1:
raise ValueError("dropout probability has to be between 0 and 1, "
"but got {}".format(p))
raise ValueError(
"dropout probability has to be between 0 and 1, "
"but got {}".format(p)
)
ctx.p = p
ctx.train = train
ctx.inplace = inplace
@@ -236,6 +262,7 @@ dropout_no_scaling = _DropoutNoScaling.apply
class _FeatureDropoutNoScaling(_DropoutNoScaling):
@staticmethod
def symbolic(input, p=0.5, train=False, inplace=False):
return None
@@ -244,7 +271,8 @@ class _FeatureDropoutNoScaling(_DropoutNoScaling):
def _make_noise(input):
return input.new().resize_(
input.size(0), input.size(1), *repeat(1,
input.dim() - 2))
input.dim() - 2)
)
feature_dropout_no_scaling = _FeatureDropoutNoScaling.apply
@@ -252,21 +280,17 @@ feature_dropout_no_scaling = _FeatureDropoutNoScaling.apply
def checkpoint_state(model=None, optimizer=None, best_prec=None, epoch=None):
return {
'epoch':
epoch,
'best_prec':
best_prec,
'model_state':
model.state_dict() if model is not None else None,
'optimizer_state':
optimizer.state_dict() if optimizer is not None else None
'epoch': epoch,
'best_prec': best_prec,
'model_state': model.state_dict() if model is not None else None,
'optimizer_state': optimizer.state_dict()
if optimizer is not None else None
}
def save_checkpoint(state,
is_best,
filename='checkpoint',
bestname='model_best'):
def save_checkpoint(
state, is_best, filename='checkpoint', bestname='model_best'
):
filename = '{}.pth.tar'.format(filename)
torch.save(state, filename)
if is_best:
@@ -325,7 +349,8 @@ def variable_size_collate(pad_val=0, use_shared_memory=True):
out = out.view(
len(batch), max_len,
*[batch[0].size(i) for i in range(1, batch[0].dim())])
*[batch[0].size(i) for i in range(1, batch[0].dim())]
)
out.fill_(pad_val)
for i in range(len(batch)):
out[i, 0:batch[i].size(0)] = batch[i]
@@ -342,8 +367,9 @@ def variable_size_collate(pad_val=0, use_shared_memory=True):
return wrapped([torch.from_numpy(b) for b in batch])
if elem.shape == (): # scalars
py_type = float if elem.dtype.name.startswith('float') else int
return _numpy_type_map[elem.dtype.name](list(
map(py_type, batch)))
return _numpy_type_map[elem.dtype.name](
list(map(py_type, batch))
)
elif isinstance(batch[0], int):
return torch.LongTensor(batch)
elif isinstance(batch[0], float):
@@ -372,19 +398,19 @@ class TrainValSplitter():
Whether or not shuffle which data goes to which split
"""
def __init__(self,
*,
numel: int,
percent_train: float,
shuffled: bool = False):
def __init__(
self, *, numel: int, percent_train: float, shuffled: bool = False
):
indicies = np.array([i for i in range(numel)])
if shuffled:
np.random.shuffle(indicies)
self.train = torch.utils.data.sampler.SubsetRandomSampler(
indicies[0:int(percent_train * numel)])
indicies[0:int(percent_train * numel)]
)
self.val = torch.utils.data.sampler.SubsetRandomSampler(
indicies[int(percent_train * numel):-1])
indicies[int(percent_train * numel):-1]
)
class CrossValSplitter():
@@ -413,7 +439,8 @@ class CrossValSplitter():
self.val = torch.utils.data.sampler.SubsetRandomSampler(self.folds[0])
self.train = torch.utils.data.sampler.SubsetRandomSampler(
np.concatenate(self.folds[1:], axis=0))
np.concatenate(self.folds[1:], axis=0)
)
self.metrics = {}
@@ -428,7 +455,8 @@ class CrossValSplitter():
assert idx >= 0 and idx < len(self)
self.val.inidicies = self.folds[idx]
self.train.inidicies = np.concatenate(
self.folds[np.arange(len(self)) != idx], axis=0)
self.folds[np.arange(len(self)) != idx], axis=0
)
def __next__(self):
self.current_v_ind += 1
@@ -454,6 +482,7 @@ class CrossValSplitter():
def set_bn_momentum_default(bn_momentum):
def fn(m):
if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
m.momentum = bn_momentum
@@ -462,14 +491,17 @@ def set_bn_momentum_default(bn_momentum):
class BNMomentumScheduler(object):
def __init__(self,
model,
bn_lambda,
last_epoch=-1,
setter=set_bn_momentum_default):
def __init__(
self, model, bn_lambda, last_epoch=-1,
setter=set_bn_momentum_default
):
if not isinstance(model, nn.Module):
raise RuntimeError("Class '{}' is not a PyTorch nn Module".format(
type(model).__name__))
raise RuntimeError(
"Class '{}' is not a PyTorch nn Module".format(
type(model).__name__
)
)
self.model = model
self.setter = setter
@@ -511,18 +543,21 @@ class Trainer(object):
Name of file to output tensorboard_logger to
"""
def __init__(self,
model,
model_fn,
optimizer,
checkpoint_name="ckpt",
best_name="best",
lr_scheduler=None,
bnm_scheduler=None,
eval_frequency=1,
log_name=None):
def __init__(
self,
model,
model_fn,
optimizer,
checkpoint_name="ckpt",
best_name="best",
lr_scheduler=None,
bnm_scheduler=None,
eval_frequency=1,
log_name=None
):
self.model, self.model_fn, self.optimizer, self.lr_scheduler, self.bnm_scheduler = (
model, model_fn, optimizer, lr_scheduler, bnm_scheduler)
model, model_fn, optimizer, lr_scheduler, bnm_scheduler
)
self.checkpoint_name, self.best_name = checkpoint_name, best_name
self.eval_frequency = eval_frequency
@@ -536,7 +571,8 @@ class Trainer(object):
@staticmethod
def _print(mode, epoch, loss, eval_dict, count):
to_print = "[{:d}] {}\tMean Loss: {:.4e}".format(
epoch, mode, loss / count)
epoch, mode, loss / count
)
for k, v in natsorted(eval_dict.items(), key=itemgetter(0)):
to_print += "\tMean {}: {:2.3f}%".format(k, stats.mean(v) * 1e2)
@@ -574,7 +610,8 @@ class Trainer(object):
for k, v in eval_res.items():
if v is not None:
tb_log.log_value(
"Training {}".format(k), 1.0 - v, step=idx)
"Training {}".format(k), 1.0 - v, step=idx
)
d_loader.dataset.randomize()
@@ -593,7 +630,8 @@ class Trainer(object):
self.optimizer.zero_grad()
_, loss, eval_res = self.model_fn(
self.model, data, eval=True, epoch=epoch)
self.model, data, eval=True, epoch=epoch
)
total_loss += loss.data[0]
count += 1
@@ -606,8 +644,7 @@ class Trainer(object):
tb_log.log_value("Eval loss", loss.data[0], step=idx)
for k, v in eval_res.items():
if v is not None:
tb_log.log_value(
"Eval {}".format(k), 1.0 - v, step=idx)
tb_log.log_value("Eval {}".format(k), 1.0 - v, step=idx)
d_loader.dataset.randomize()
@@ -615,12 +652,14 @@ class Trainer(object):
return total_loss / count, eval_dict
def train(self,
start_epoch,
n_epochs,
train_loader,
test_loader=None,
best_loss=0.0):
def train(
self,
start_epoch,
n_epochs,
train_loader,
test_loader=None,
best_loss=0.0
):
r"""
Call to begin training the model
@@ -649,10 +688,12 @@ class Trainer(object):
is_best = val_loss < best_loss
best_loss = min(best_loss, val_loss)
save_checkpoint(
checkpoint_state(self.model, self.optimizer, val_loss,
epoch),
checkpoint_state(
self.model, self.optimizer, val_loss, epoch
),
is_best,
filename=self.checkpoint_name,
bestname=self.best_name)
bestname=self.best_name
)
return best_loss