mirror of
https://github.com/wassname/Pointnet2_PyTorch.git
synced 2026-07-25 13:00:37 +08:00
Some faster/better kernels. Tensors with points are now kept in (b, c, ...) format as this is easier for pytorch
This commit is contained in:
+18
-2
@@ -4,7 +4,7 @@
|
||||
|
||||
extern THCState *state;
|
||||
|
||||
int gather_points_wrapper(int b, int n, int c, int npoints,
|
||||
int gather_points_wrapper(int b, int c, int n, int npoints,
|
||||
THCudaTensor *points_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *out_tensor) {
|
||||
@@ -15,7 +15,23 @@ int gather_points_wrapper(int b, int n, int c, int npoints,
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
gather_points_kernel_wrapper(b, n, c, npoints, points, idx, out, stream);
|
||||
gather_points_kernel_wrapper(b, c, n, npoints, points, idx, out, stream);
|
||||
return 1;
|
||||
}
|
||||
|
||||
int gather_points_grad_wrapper(int b, int c, int n, int npoints,
|
||||
THCudaTensor *grad_out_tensor,
|
||||
THCudaIntTensor *idx_tensor,
|
||||
THCudaTensor *grad_points_tensor) {
|
||||
|
||||
const float *grad_out = THCudaTensor_data(state, grad_out_tensor);
|
||||
const int *idx = THCudaIntTensor_data(state, idx_tensor);
|
||||
float *grad_points = THCudaTensor_data(state, grad_points_tensor);
|
||||
|
||||
cudaStream_t stream = THCState_getCurrentStream(state);
|
||||
|
||||
gather_points_grad_kernel_wrapper(b, c, n, npoints, grad_out, idx,
|
||||
grad_points, stream);
|
||||
return 1;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user