Initial commit

Author: chrdiller

.gitignore

@@ -0,0 +1,3 @@
+.DS_Store
+._*
+

LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [year] [fullname]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,20 @@
+# Chamfer Distance for pyTorch
+
+This is an implementation of the Chamfer Distance as a module for pyTorch. It is written as a custom C++/CUDA extension.
+
+As it is using pyTorch's [JIT compilation](https://pytorch.org/tutorials/advanced/cpp_extension.html), there are no additional prerequisite steps that have to be taken. Simply import the module as shown below; CUDA and C++ code will be compiled on the first run.
+
+### Usage
+```python
+from chamfer_distance import ChamferDistance
+chamfer_dist = ChamferDistance()
+
+#...
+# points and points_reconstructed are n_points x 3 matrices
+
+dist1, dist2 = chamfer_dist(points, points_reconstructed)
+loss = (torch.mean(dist1)) + (torch.mean(dist2))
+
+
+#...
+```

chamfer_distance/__init__.py

@@ -0,0 +1 @@
+from .chamfer_distance import ChamferDistance

chamfer_distance/chamfer_distance.cpp

@@ -0,0 +1,185 @@
+#include <torch/torch.h>
+
+// CUDA forward declarations
+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)");
+}