|
#include <torch/torch.h> |
|
|
|
|
|
int ChamferDistanceKernelLauncher( |
|
const int b, const int n, |
|
const float* xyz, |
|
const int m, |
|
const float* xyz2, |
|
float* result, |
|
int* result_i, |
|
float* result2, |
|
int* result2_i); |
|
|
|
int ChamferDistanceGradKernelLauncher( |
|
const int b, const int n, |
|
const float* xyz1, |
|
const int m, |
|
const float* xyz2, |
|
const float* grad_dist1, |
|
const int* idx1, |
|
const float* grad_dist2, |
|
const int* idx2, |
|
float* grad_xyz1, |
|
float* grad_xyz2); |
|
|
|
|
|
void chamfer_distance_forward_cuda( |
|
const at::Tensor xyz1, |
|
const at::Tensor xyz2, |
|
const at::Tensor dist1, |
|
const at::Tensor dist2, |
|
const at::Tensor idx1, |
|
const at::Tensor idx2) |
|
{ |
|
ChamferDistanceKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data<float>(), |
|
xyz2.size(1), xyz2.data<float>(), |
|
dist1.data<float>(), idx1.data<int>(), |
|
dist2.data<float>(), idx2.data<int>()); |
|
} |
|
|
|
void chamfer_distance_backward_cuda( |
|
const at::Tensor xyz1, |
|
const at::Tensor xyz2, |
|
at::Tensor gradxyz1, |
|
at::Tensor gradxyz2, |
|
at::Tensor graddist1, |
|
at::Tensor graddist2, |
|
at::Tensor idx1, |
|
at::Tensor idx2) |
|
{ |
|
ChamferDistanceGradKernelLauncher(xyz1.size(0), xyz1.size(1), xyz1.data<float>(), |
|
xyz2.size(1), xyz2.data<float>(), |
|
graddist1.data<float>(), idx1.data<int>(), |
|
graddist2.data<float>(), idx2.data<int>(), |
|
gradxyz1.data<float>(), gradxyz2.data<float>()); |
|
} |
|
|
|
|
|
void nnsearch( |
|
const int b, const int n, const int m, |
|
const float* xyz1, |
|
const float* xyz2, |
|
float* dist, |
|
int* idx) |
|
{ |
|
for (int i = 0; i < b; i++) { |
|
for (int j = 0; j < n; j++) { |
|
const float x1 = xyz1[(i*n+j)*3+0]; |
|
const float y1 = xyz1[(i*n+j)*3+1]; |
|
const float z1 = xyz1[(i*n+j)*3+2]; |
|
double best = 0; |
|
int besti = 0; |
|
for (int k = 0; k < m; k++) { |
|
const float x2 = xyz2[(i*m+k)*3+0] - x1; |
|
const float y2 = xyz2[(i*m+k)*3+1] - y1; |
|
const float z2 = xyz2[(i*m+k)*3+2] - z1; |
|
const double d=x2*x2+y2*y2+z2*z2; |
|
if (k==0 || d < best){ |
|
best = d; |
|
besti = k; |
|
} |
|
} |
|
dist[i*n+j] = best; |
|
idx[i*n+j] = besti; |
|
} |
|
} |
|
} |
|
|
|
|
|
void chamfer_distance_forward( |
|
const at::Tensor xyz1, |
|
const at::Tensor xyz2, |
|
const at::Tensor dist1, |
|
const at::Tensor dist2, |
|
const at::Tensor idx1, |
|
const at::Tensor idx2) |
|
{ |
|
const int batchsize = xyz1.size(0); |
|
const int n = xyz1.size(1); |
|
const int m = xyz2.size(1); |
|
|
|
const float* xyz1_data = xyz1.data<float>(); |
|
const float* xyz2_data = xyz2.data<float>(); |
|
float* dist1_data = dist1.data<float>(); |
|
float* dist2_data = dist2.data<float>(); |
|
int* idx1_data = idx1.data<int>(); |
|
int* idx2_data = idx2.data<int>(); |
|
|
|
nnsearch(batchsize, n, m, xyz1_data, xyz2_data, dist1_data, idx1_data); |
|
nnsearch(batchsize, m, n, xyz2_data, xyz1_data, dist2_data, idx2_data); |
|
} |
|
|
|
|
|
void chamfer_distance_backward( |
|
const at::Tensor xyz1, |
|
const at::Tensor xyz2, |
|
at::Tensor gradxyz1, |
|
at::Tensor gradxyz2, |
|
at::Tensor graddist1, |
|
at::Tensor graddist2, |
|
at::Tensor idx1, |
|
at::Tensor idx2) |
|
{ |
|
const int b = xyz1.size(0); |
|
const int n = xyz1.size(1); |
|
const int m = xyz2.size(1); |
|
|
|
const float* xyz1_data = xyz1.data<float>(); |
|
const float* xyz2_data = xyz2.data<float>(); |
|
float* gradxyz1_data = gradxyz1.data<float>(); |
|
float* gradxyz2_data = gradxyz2.data<float>(); |
|
float* graddist1_data = graddist1.data<float>(); |
|
float* graddist2_data = graddist2.data<float>(); |
|
const int* idx1_data = idx1.data<int>(); |
|
const int* idx2_data = idx2.data<int>(); |
|
|
|
for (int i = 0; i < b*n*3; i++) |
|
gradxyz1_data[i] = 0; |
|
for (int i = 0; i < b*m*3; i++) |
|
gradxyz2_data[i] = 0; |
|
for (int i = 0;i < b; i++) { |
|
for (int j = 0; j < n; j++) { |
|
const float x1 = xyz1_data[(i*n+j)*3+0]; |
|
const float y1 = xyz1_data[(i*n+j)*3+1]; |
|
const float z1 = xyz1_data[(i*n+j)*3+2]; |
|
const int j2 = idx1_data[i*n+j]; |
|
|
|
const float x2 = xyz2_data[(i*m+j2)*3+0]; |
|
const float y2 = xyz2_data[(i*m+j2)*3+1]; |
|
const float z2 = xyz2_data[(i*m+j2)*3+2]; |
|
const float g = graddist1_data[i*n+j]*2; |
|
|
|
gradxyz1_data[(i*n+j)*3+0] += g*(x1-x2); |
|
gradxyz1_data[(i*n+j)*3+1] += g*(y1-y2); |
|
gradxyz1_data[(i*n+j)*3+2] += g*(z1-z2); |
|
gradxyz2_data[(i*m+j2)*3+0] -= (g*(x1-x2)); |
|
gradxyz2_data[(i*m+j2)*3+1] -= (g*(y1-y2)); |
|
gradxyz2_data[(i*m+j2)*3+2] -= (g*(z1-z2)); |
|
} |
|
for (int j = 0; j < m; j++) { |
|
const float x1 = xyz2_data[(i*m+j)*3+0]; |
|
const float y1 = xyz2_data[(i*m+j)*3+1]; |
|
const float z1 = xyz2_data[(i*m+j)*3+2]; |
|
const int j2 = idx2_data[i*m+j]; |
|
const float x2 = xyz1_data[(i*n+j2)*3+0]; |
|
const float y2 = xyz1_data[(i*n+j2)*3+1]; |
|
const float z2 = xyz1_data[(i*n+j2)*3+2]; |
|
const float g = graddist2_data[i*m+j]*2; |
|
gradxyz2_data[(i*m+j)*3+0] += g*(x1-x2); |
|
gradxyz2_data[(i*m+j)*3+1] += g*(y1-y2); |
|
gradxyz2_data[(i*m+j)*3+2] += g*(z1-z2); |
|
gradxyz1_data[(i*n+j2)*3+0] -= (g*(x1-x2)); |
|
gradxyz1_data[(i*n+j2)*3+1] -= (g*(y1-y2)); |
|
gradxyz1_data[(i*n+j2)*3+2] -= (g*(z1-z2)); |
|
} |
|
} |
|
} |
|
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { |
|
m.def("forward", &chamfer_distance_forward, "ChamferDistance forward"); |
|
m.def("forward_cuda", &chamfer_distance_forward_cuda, "ChamferDistance forward (CUDA)"); |
|
m.def("backward", &chamfer_distance_backward, "ChamferDistance backward"); |
|
m.def("backward_cuda", &chamfer_distance_backward_cuda, "ChamferDistance backward (CUDA)"); |
|
} |
|
|
