initial commit

.gitignore

@@ -0,0 +1,24 @@
+__pycache__
+*.pth*
+.autoenv*
+runs
+build
+checkpoints
+*.prof
+.lvimrc
+.vimtags
+.ccls
+.ccls-cache/
+dist/
+pointnet2.egg-info/
+*.zip
+*.so
+.tox/
+.mypy_cache
+**/*.pyc
+
+
+pointnet2/data/modelnet40_normal_resampled/
+pointnet2/data/modelnet40_normal_resampled_cache/
+pointnet2/data/modelnet40_ply_hdf5_2048/
+outputs/

PytorchEMD/.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+build
+dist
+emd_ext.egg-info
+*.so

PytorchEMD/README.md

@@ -0,0 +1,31 @@
+# PyTorch Wrapper for Point-cloud Earth-Mover-Distance (EMD)
+
+## Dependency
+
+The code has been tested on Ubuntu 16.04, PyTorch 1.1.0, CUDA 9.0.
+
+## Usage
+
+First compile using
+
+        python setup.py install
+
+Then, copy the lib file out to the main directory,
+
+        cp build/lib.linux-x86_64-3.6/emd_cuda.cpython-36m-x86_64-linux-gnu.so .
+
+Then, you can use it by simply
+
+        from emd import earth_mover_distance
+        d = earth_mover_distance(p1, p2, transpose=False)  # p1: B x N1 x 3, p2: B x N2 x 3
+
+Check `test_emd_loss.py` for example.
+
+## Author
+
+The cuda code is originally written by Haoqiang Fan. The PyTorch wrapper is written by Kaichun Mo. Also, Jiayuan Gu provided helps.
+
+## License
+
+MIT
+

PytorchEMD/cuda/emd.cpp

@@ -0,0 +1,29 @@
+#ifndef _EMD
+#define _EMD
+
+#include <vector>
+#include <torch/extension.h>
+
+//CUDA declarations
+at::Tensor ApproxMatchForward(
+    const at::Tensor xyz1,
+    const at::Tensor xyz2);
+
+at::Tensor MatchCostForward(
+    const at::Tensor xyz1,
+    const at::Tensor xyz2,
+    const at::Tensor match);
+
+std::vector<at::Tensor> MatchCostBackward(
+    const at::Tensor grad_cost,
+    const at::Tensor xyz1,
+    const at::Tensor xyz2,
+    const at::Tensor match);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("approxmatch_forward", &ApproxMatchForward,"ApproxMatch forward (CUDA)");
+  m.def("matchcost_forward", &MatchCostForward,"MatchCost forward (CUDA)");
+  m.def("matchcost_backward", &MatchCostBackward,"MatchCost backward (CUDA)");
+}
+
+#endif