wassname/Pointnet2_PyTorch
1
0
Fork 0
mirror of https://github.com/wassname/Pointnet2_PyTorch.git synced 2026-07-20 12:10:53 +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/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;
}