openpoints 0.1.0__tar.gz

openpoints-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Guocheng Qian.
+
+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.

openpoints-0.1.0/MANIFEST.in

@@ -0,0 +1,6 @@
+include README.md
+include LICENSE
+recursive-include cpp *.cpp *.cu *.cxx *.h *.hpp *.pyx *.md *.py
+recursive-include dataset/semantic_kitti *.yaml *.txt *.cpp *.cxx *.h *.hpp *.pyx *.sh
+include dataset/matterport3d/category_mapping.tsv
+global-exclude *.py[cod] __pycache__/* *.so .DS_Store

openpoints-0.1.0/PKG-INFO

@@ -0,0 +1,150 @@
+Metadata-Version: 2.4
+Name: openpoints
+Version: 0.1.0
+Summary: OpenPoints: point cloud understanding models, layers, datasets, transforms, optimizers, and training utilities.
+Author: Guocheng Qian
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/guochengqian/openpoints
+Project-URL: Repository, https://github.com/guochengqian/openpoints
+Project-URL: Issues, https://github.com/guochengqian/PointNeXt/issues
+Keywords: point cloud,deep learning,pytorch,3d vision,PointNeXt
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.10
+Requires-Dist: numpy>=1.20
+Requires-Dist: scipy>=1.7
+Requires-Dist: PyYAML>=5.4
+Requires-Dist: multimethod>=1.8
+Requires-Dist: termcolor>=1.1
+Requires-Dist: shortuuid>=1.0
+Requires-Dist: scikit-learn>=1.0
+Requires-Dist: tqdm>=4
+Requires-Dist: easydict>=1.9
+Requires-Dist: h5py>=3
+Requires-Dist: pandas>=1.3
+Requires-Dist: torchvision>=0.11
+Provides-Extra: data
+Requires-Dist: easydict>=1.9; extra == "data"
+Requires-Dist: h5py>=3; extra == "data"
+Requires-Dist: pandas>=1.3; extra == "data"
+Requires-Dist: torchvision>=0.11; extra == "data"
+Provides-Extra: viz
+Requires-Dist: matplotlib>=3.5; extra == "viz"
+Requires-Dist: pyvista>=0.38; extra == "viz"
+Provides-Extra: graph
+Requires-Dist: numba>=0.56; extra == "graph"
+Requires-Dist: ogb>=1.3; extra == "graph"
+Requires-Dist: atom3d>=0.2; extra == "graph"
+Requires-Dist: fast-pytorch-kmeans>=0.1; extra == "graph"
+Provides-Extra: tfds
+Requires-Dist: Pillow>=8; extra == "tfds"
+Requires-Dist: tensorflow>=2; extra == "tfds"
+Requires-Dist: tensorflow-datasets>=4; extra == "tfds"
+Provides-Extra: wandb
+Requires-Dist: wandb>=0.13; extra == "wandb"
+Provides-Extra: dev
+Requires-Dist: build>=1; extra == "dev"
+Requires-Dist: twine>=4; extra == "dev"
+Requires-Dist: pytest>=7; extra == "dev"
+Provides-Extra: all
+Requires-Dist: easydict>=1.9; extra == "all"
+Requires-Dist: h5py>=3; extra == "all"
+Requires-Dist: pandas>=1.3; extra == "all"
+Requires-Dist: torchvision>=0.11; extra == "all"
+Requires-Dist: matplotlib>=3.5; extra == "all"
+Requires-Dist: pyvista>=0.38; extra == "all"
+Requires-Dist: numba>=0.56; extra == "all"
+Requires-Dist: ogb>=1.3; extra == "all"
+Requires-Dist: atom3d>=0.2; extra == "all"
+Requires-Dist: fast-pytorch-kmeans>=0.1; extra == "all"
+Requires-Dist: Pillow>=8; extra == "all"
+Requires-Dist: tensorflow>=2; extra == "all"
+Requires-Dist: tensorflow-datasets>=4; extra == "all"
+Requires-Dist: wandb>=0.13; extra == "all"
+Dynamic: license-file
+
+# OpenPoints
+
+OpenPoints is a library built for fairly benchmarking and easily reproducing point-based methods for point cloud understanding. It is born in the course of [PointNeXt](https://github.com/guochengqian/PointNeXt) project and is used as an engine therein.
+
+**For any question related to OpenPoints, please open an issue in [PointNeXt](https://github.com/guochengqian/PointNeXt) repo.**
+
+OpenPoints currently supports reproducing the following models:
+- PointNet
+- DGCNN
+- DeepGCN
+- PointNet++
+- ASSANet
+- PointMLP
+- PointNeXt
+- Pix4Point
+- PointVector
+
+
+
+## Features
+
+1. **Extensibility**: supports many representative networks for point cloud understanding, such as *PointNet, DGCNN, DeepGCN, PointNet++, ASSANet, PointMLP*, and our ***PointNeXt***. More networks can be built easily based on our framework since **OpenPoints support a wide range of basic operations including graph convolutions, self-attention, farthest point sampling, ball query, *e.t.c***.
+
+2. **Ease of Use**: *Build* model, optimizer, scheduler, loss function,  and data loader *easily from cfg*. Train and validate different models on various tasks by simply changing the `cfg\*\*.yaml` file. 
+
+      ```
+   model = build_model_from_cfg(cfg.model)
+   criterion = build_criterion_from_cfg(cfg.criterion_args)
+   ```
+
+
+
+## Installation
+
+OpenPoints can be installed as the `openpoints` Python package:
+
+```bash
+pip install openpoints
+```
+
+The PyPI package installs the Python library (`import openpoints`) and the files needed by the datasets and configs. CUDA/C++ operators such as `pointnet2_batch_cuda`, `pointops_cuda`, `chamfer`, and `emd_cuda` are still built from a source checkout or source distribution for now because PyTorch/CUDA wheels must match the user's Python, PyTorch, CUDA, and platform versions.
+
+For full training/evaluation with CUDA ops, install from source after installing PyTorch:
+
+```bash
+git clone --recursive https://github.com/guochengqian/openpoints.git
+cd openpoints
+pip install -e .[data,viz,wandb]
+cd cpp/pointnet2_batch && python setup.py install && cd ../..
+cd cpp/pointops && python setup.py install && cd ../..
+cd cpp/chamfer_dist && python setup.py install && cd ../..
+cd cpp/emd && python setup.py install && cd ../..
+```
+
+If the `openpoints` name is unavailable on a package index mirror, the same code can be published as `openpoints-torch`; the import name remains `openpoints`.
+
+## Usage
+
+OpenPoints serves as the engine for PointNeXt. Please refer to [PointNeXt](https://github.com/guochengqian/PointNeXt) for complete training, evaluation, and model-zoo examples.
+
+
+
+## Citation
+
+If you use this library, please kindly acknowledge our work:
+```tex
+@Article{qian2022pointnext,
+  author  = {Qian, Guocheng and Li, Yuchen and Peng, Houwen and Mai, Jinjie and Hammoud, Hasan and Elhoseiny, Mohamed and Ghanem, Bernard},
+  title   = {PointNeXt: Revisiting PointNet++ with Improved Training and Scaling Strategies},
+  journal = {arXiv:2206.04670},
+  year    = {2022},
+}
+```
+

openpoints-0.1.0/README.md

@@ -0,0 +1,74 @@
+# OpenPoints
+
+OpenPoints is a library built for fairly benchmarking and easily reproducing point-based methods for point cloud understanding. It is born in the course of [PointNeXt](https://github.com/guochengqian/PointNeXt) project and is used as an engine therein.
+
+**For any question related to OpenPoints, please open an issue in [PointNeXt](https://github.com/guochengqian/PointNeXt) repo.**
+
+OpenPoints currently supports reproducing the following models:
+- PointNet
+- DGCNN
+- DeepGCN
+- PointNet++
+- ASSANet
+- PointMLP
+- PointNeXt
+- Pix4Point
+- PointVector
+
+
+
+## Features
+
+1. **Extensibility**: supports many representative networks for point cloud understanding, such as *PointNet, DGCNN, DeepGCN, PointNet++, ASSANet, PointMLP*, and our ***PointNeXt***. More networks can be built easily based on our framework since **OpenPoints support a wide range of basic operations including graph convolutions, self-attention, farthest point sampling, ball query, *e.t.c***.
+
+2. **Ease of Use**: *Build* model, optimizer, scheduler, loss function,  and data loader *easily from cfg*. Train and validate different models on various tasks by simply changing the `cfg\*\*.yaml` file. 
+
+      ```
+   model = build_model_from_cfg(cfg.model)
+   criterion = build_criterion_from_cfg(cfg.criterion_args)
+   ```
+
+
+
+## Installation
+
+OpenPoints can be installed as the `openpoints` Python package:
+
+```bash
+pip install openpoints
+```
+
+The PyPI package installs the Python library (`import openpoints`) and the files needed by the datasets and configs. CUDA/C++ operators such as `pointnet2_batch_cuda`, `pointops_cuda`, `chamfer`, and `emd_cuda` are still built from a source checkout or source distribution for now because PyTorch/CUDA wheels must match the user's Python, PyTorch, CUDA, and platform versions.
+
+For full training/evaluation with CUDA ops, install from source after installing PyTorch:
+
+```bash
+git clone --recursive https://github.com/guochengqian/openpoints.git
+cd openpoints
+pip install -e .[data,viz,wandb]
+cd cpp/pointnet2_batch && python setup.py install && cd ../..
+cd cpp/pointops && python setup.py install && cd ../..
+cd cpp/chamfer_dist && python setup.py install && cd ../..
+cd cpp/emd && python setup.py install && cd ../..
+```
+
+If the `openpoints` name is unavailable on a package index mirror, the same code can be published as `openpoints-torch`; the import name remains `openpoints`.
+
+## Usage
+
+OpenPoints serves as the engine for PointNeXt. Please refer to [PointNeXt](https://github.com/guochengqian/PointNeXt) for complete training, evaluation, and model-zoo examples.
+
+
+
+## Citation
+
+If you use this library, please kindly acknowledge our work:
+```tex
+@Article{qian2022pointnext,
+  author  = {Qian, Guocheng and Li, Yuchen and Peng, Houwen and Mai, Jinjie and Hammoud, Hasan and Elhoseiny, Mohamed and Ghanem, Bernard},
+  title   = {PointNeXt: Revisiting PointNet++ with Improved Training and Scaling Strategies},
+  journal = {arXiv:2206.04670},
+  year    = {2022},
+}
+```
+

openpoints-0.1.0/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "0.1.0"
+
+from .transforms import *

openpoints-0.1.0/cpp/__init__.py

@@ -0,0 +1,6 @@
+"""
+Author: PointNeXt
+
+"""
+
+from .pointnet2_batch import pointnet2_cuda

openpoints-0.1.0/cpp/chamfer_dist/__init__.py

@@ -0,0 +1,110 @@
+# -*- coding: utf-8 -*-
+# @Author: Thibault GROUEIX
+# @Date:   2019-08-07 20:54:24
+# @Last Modified by:   Haozhe Xie
+# @Last Modified time: 2019-12-18 15:06:25
+# @Email:  cshzxie@gmail.com
+"""Chamfer distance CUDA extension loader.
+
+The PyPI package ships Python sources and keeps compiled CUDA extensions
+optional. Importing this module should not fail on machines where the
+``chamfer`` extension has not been built yet; users only need the extension
+when they instantiate/use Chamfer distance losses.
+"""
+
+import torch
+
+try:
+    import chamfer as _chamfer
+except ImportError as exc:
+    _chamfer = None
+    _chamfer_import_error = exc
+else:
+    _chamfer_import_error = None
+
+
+def _require_chamfer():
+    if _chamfer is None:
+        raise ImportError(
+            "The `chamfer` CUDA extension is not installed. Build OpenPoints "
+            "CUDA ops from source first, e.g. `cd cpp/chamfer_dist && "
+            "python setup.py install`, or install a wheel that includes the "
+            "compiled extension."
+        ) from _chamfer_import_error
+    return _chamfer
+
+
+class ChamferFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, xyz1, xyz2):
+        chamfer = _require_chamfer()
+        dist1, dist2, idx1, idx2 = chamfer.forward(xyz1, xyz2)
+        ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
+
+        return dist1, dist2
+
+    @staticmethod
+    def backward(ctx, grad_dist1, grad_dist2):
+        chamfer = _require_chamfer()
+        xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
+        grad_xyz1, grad_xyz2 = chamfer.backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2)
+        return grad_xyz1, grad_xyz2
+
+
+class ChamferDistanceL2(torch.nn.Module):
+    f''' Chamder Distance L2
+    '''
+    def __init__(self, ignore_zeros=False):
+        super().__init__()
+        self.ignore_zeros = ignore_zeros
+
+    def forward(self, xyz1, xyz2):
+        batch_size = xyz1.size(0)
+        if batch_size == 1 and self.ignore_zeros:
+            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+        return torch.mean(dist1) + torch.mean(dist2)
+
+class ChamferDistanceL2_split(torch.nn.Module):
+    f''' Chamder Distance L2
+    '''
+    def __init__(self, ignore_zeros=False):
+        super().__init__()
+        self.ignore_zeros = ignore_zeros
+
+    def forward(self, xyz1, xyz2):
+        batch_size = xyz1.size(0)
+        if batch_size == 1 and self.ignore_zeros:
+            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+        return torch.mean(dist1), torch.mean(dist2)
+
+class ChamferDistanceL1(torch.nn.Module):
+    f''' Chamder Distance L1
+    '''
+    def __init__(self, ignore_zeros=False):
+        super().__init__()
+        self.ignore_zeros = ignore_zeros
+
+    def forward(self, xyz1, xyz2):
+        batch_size = xyz1.size(0)
+        if batch_size == 1 and self.ignore_zeros:
+            non_zeros1 = torch.sum(xyz1, dim=2).ne(0)
+            non_zeros2 = torch.sum(xyz2, dim=2).ne(0)
+            xyz1 = xyz1[non_zeros1].unsqueeze(dim=0)
+            xyz2 = xyz2[non_zeros2].unsqueeze(dim=0)
+
+        dist1, dist2 = ChamferFunction.apply(xyz1, xyz2)
+        # import pdb
+        # pdb.set_trace()
+        dist1 = torch.sqrt(dist1)
+        dist2 = torch.sqrt(dist2)
+        return (torch.mean(dist1) + torch.mean(dist2))/2

openpoints-0.1.0/cpp/chamfer_dist/chamfer.cu

@@ -0,0 +1,229 @@
+/*
+ * @Author: Haozhe Xie
+ * @Date:   2019-08-07 20:54:24
+ * @Last Modified by:   Haozhe Xie
+ * @Last Modified time: 2020-06-17 14:58:55
+ * @Email:  cshzxie@gmail.com
+ */
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <vector>
+
+__global__ void chamfer_dist_kernel(int batch_size,
+                                    int n,
+                                    const float* xyz1,
+                                    int m,
+                                    const float* xyz2,
+                                    float* dist,
+                                    int* indexes) {
+  const int batch = 512;
+  __shared__ float buf[batch * 3];
+  for (int i = blockIdx.x; i < batch_size; i += gridDim.x) {
+    for (int k2 = 0; k2 < m; k2 += batch) {
+      int end_k = min(m, k2 + batch) - k2;
+      for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) {
+        buf[j] = xyz2[(i * m + k2) * 3 + j];
+      }
+      __syncthreads();
+      for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n;
+           j += blockDim.x * gridDim.y) {
+        float x1            = xyz1[(i * n + j) * 3 + 0];
+        float y1            = xyz1[(i * n + j) * 3 + 1];
+        float z1            = xyz1[(i * n + j) * 3 + 2];
+        float best_dist     = 0;
+        int best_dist_index = 0;
+        int end_ka          = end_k - (end_k & 3);
+        if (end_ka == batch) {
+          for (int k = 0; k < batch; k += 4) {
+            {
+              float x2   = buf[k * 3 + 0] - x1;
+              float y2   = buf[k * 3 + 1] - y1;
+              float z2   = buf[k * 3 + 2] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+
+              if (k == 0 || dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2;
+              }
+            }
+            {
+              float x2   = buf[k * 3 + 3] - x1;
+              float y2   = buf[k * 3 + 4] - y1;
+              float z2   = buf[k * 3 + 5] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2 + 1;
+              }
+            }
+            {
+              float x2   = buf[k * 3 + 6] - x1;
+              float y2   = buf[k * 3 + 7] - y1;
+              float z2   = buf[k * 3 + 8] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2 + 2;
+              }
+            }
+            {
+              float x2   = buf[k * 3 + 9] - x1;
+              float y2   = buf[k * 3 + 10] - y1;
+              float z2   = buf[k * 3 + 11] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2 + 3;
+              }
+            }
+          }
+        } else {
+          for (int k = 0; k < end_ka; k += 4) {
+            {
+              float x2   = buf[k * 3 + 0] - x1;
+              float y2   = buf[k * 3 + 1] - y1;
+              float z2   = buf[k * 3 + 2] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (k == 0 || dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2;
+              }
+            }
+            {
+              float x2   = buf[k * 3 + 3] - x1;
+              float y2   = buf[k * 3 + 4] - y1;
+              float z2   = buf[k * 3 + 5] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2 + 1;
+              }
+            }
+            {
+              float x2   = buf[k * 3 + 6] - x1;
+              float y2   = buf[k * 3 + 7] - y1;
+              float z2   = buf[k * 3 + 8] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2 + 2;
+              }
+            }
+            {
+              float x2   = buf[k * 3 + 9] - x1;
+              float y2   = buf[k * 3 + 10] - y1;
+              float z2   = buf[k * 3 + 11] - z1;
+              float dist = x2 * x2 + y2 * y2 + z2 * z2;
+              if (dist < best_dist) {
+                best_dist       = dist;
+                best_dist_index = k + k2 + 3;
+              }
+            }
+          }
+        }
+        for (int k = end_ka; k < end_k; k++) {
+          float x2   = buf[k * 3 + 0] - x1;
+          float y2   = buf[k * 3 + 1] - y1;
+          float z2   = buf[k * 3 + 2] - z1;
+          float dist = x2 * x2 + y2 * y2 + z2 * z2;
+          if (k == 0 || dist < best_dist) {
+            best_dist       = dist;
+            best_dist_index = k + k2;
+          }
+        }
+        if (k2 == 0 || dist[(i * n + j)] > best_dist) {
+          dist[(i * n + j)]    = best_dist;
+          indexes[(i * n + j)] = best_dist_index;
+        }
+      }
+      __syncthreads();
+    }
+  }
+}
+
+std::vector<torch::Tensor> chamfer_cuda_forward(torch::Tensor xyz1,
+                                                torch::Tensor xyz2) {
+  const int batch_size = xyz1.size(0);
+  const int n          = xyz1.size(1);  // num_points point cloud A
+  const int m          = xyz2.size(1);  // num_points point cloud B
+  torch::Tensor dist1 =
+    torch::zeros({batch_size, n}, torch::CUDA(torch::kFloat));
+  torch::Tensor dist2 =
+    torch::zeros({batch_size, m}, torch::CUDA(torch::kFloat));
+  torch::Tensor idx1 = torch::zeros({batch_size, n}, torch::CUDA(torch::kInt));
+  torch::Tensor idx2 = torch::zeros({batch_size, m}, torch::CUDA(torch::kInt));
+
+  chamfer_dist_kernel<<<dim3(32, 16, 1), 512>>>(
+    batch_size, n, xyz1.data_ptr<float>(), m, xyz2.data_ptr<float>(),
+    dist1.data_ptr<float>(), idx1.data_ptr<int>());
+  chamfer_dist_kernel<<<dim3(32, 16, 1), 512>>>(
+    batch_size, m, xyz2.data_ptr<float>(), n, xyz1.data_ptr<float>(),
+    dist2.data_ptr<float>(), idx2.data_ptr<int>());
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("Error in chamfer_cuda_forward: %s\n", cudaGetErrorString(err));
+  }
+  return {dist1, dist2, idx1, idx2};
+}
+
+__global__ void chamfer_dist_grad_kernel(int b,
+                                         int n,
+                                         const float* xyz1,
+                                         int m,
+                                         const float* xyz2,
+                                         const float* grad_dist1,
+                                         const int* idx1,
+                                         float* grad_xyz1,
+                                         float* grad_xyz2) {
+  for (int i = blockIdx.x; i < b; i += gridDim.x) {
+    for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n;
+         j += blockDim.x * gridDim.y) {
+      float x1 = xyz1[(i * n + j) * 3 + 0];
+      float y1 = xyz1[(i * n + j) * 3 + 1];
+      float z1 = xyz1[(i * n + j) * 3 + 2];
+      int j2   = idx1[i * n + j];
+      float x2 = xyz2[(i * m + j2) * 3 + 0];
+      float y2 = xyz2[(i * m + j2) * 3 + 1];
+      float z2 = xyz2[(i * m + j2) * 3 + 2];
+      float g  = grad_dist1[i * n + j] * 2;
+      atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 0]), g * (x1 - x2));
+      atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 1]), g * (y1 - y2));
+      atomicAdd(&(grad_xyz1[(i * n + j) * 3 + 2]), g * (z1 - z2));
+      atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 0]), -(g * (x1 - x2)));
+      atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 1]), -(g * (y1 - y2)));
+      atomicAdd(&(grad_xyz2[(i * m + j2) * 3 + 2]), -(g * (z1 - z2)));
+    }
+  }
+}
+
+std::vector<torch::Tensor> chamfer_cuda_backward(torch::Tensor xyz1,
+                                                 torch::Tensor xyz2,
+                                                 torch::Tensor idx1,
+                                                 torch::Tensor idx2,
+                                                 torch::Tensor grad_dist1,
+                                                 torch::Tensor grad_dist2) {
+  const int batch_size    = xyz1.size(0);
+  const int n             = xyz1.size(1);  // num_points point cloud A
+  const int m             = xyz2.size(1);  // num_points point cloud B
+  torch::Tensor grad_xyz1 = torch::zeros_like(xyz1, torch::CUDA(torch::kFloat));
+  torch::Tensor grad_xyz2 = torch::zeros_like(xyz2, torch::CUDA(torch::kFloat));
+
+  chamfer_dist_grad_kernel<<<dim3(1, 16, 1), 256>>>(
+    batch_size, n, xyz1.data_ptr<float>(), m, xyz2.data_ptr<float>(),
+    grad_dist1.data_ptr<float>(), idx1.data_ptr<int>(),
+    grad_xyz1.data_ptr<float>(), grad_xyz2.data_ptr<float>());
+  chamfer_dist_grad_kernel<<<dim3(1, 16, 1), 256>>>(
+    batch_size, m, xyz2.data_ptr<float>(), n, xyz1.data_ptr<float>(),
+    grad_dist2.data_ptr<float>(), idx2.data_ptr<int>(),
+    grad_xyz2.data_ptr<float>(), grad_xyz1.data_ptr<float>());
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess) {
+    printf("Error in chamfer_cuda_backward: %s\n", cudaGetErrorString(err));
+  }
+  return {grad_xyz1, grad_xyz2};
+}

openpoints-0.1.0/cpp/chamfer_dist/chamfer_cuda.cpp

@@ -0,0 +1,39 @@
+/*
+ * @Author: Haozhe Xie
+ * @Date:   2019-08-07 20:54:24
+ * @Last Modified by:   Haozhe Xie
+ * @Last Modified time: 2019-12-10 10:33:50
+ * @Email:  cshzxie@gmail.com
+ */
+
+#include <torch/extension.h>
+#include <vector>
+
+std::vector<torch::Tensor> chamfer_cuda_forward(torch::Tensor xyz1,
+                                                torch::Tensor xyz2);
+
+std::vector<torch::Tensor> chamfer_cuda_backward(torch::Tensor xyz1,
+                                                 torch::Tensor xyz2,
+                                                 torch::Tensor idx1,
+                                                 torch::Tensor idx2,
+                                                 torch::Tensor grad_dist1,
+                                                 torch::Tensor grad_dist2);
+
+std::vector<torch::Tensor> chamfer_forward(torch::Tensor xyz1,
+                                           torch::Tensor xyz2) {
+  return chamfer_cuda_forward(xyz1, xyz2);
+}
+
+std::vector<torch::Tensor> chamfer_backward(torch::Tensor xyz1,
+                                            torch::Tensor xyz2,
+                                            torch::Tensor idx1,
+                                            torch::Tensor idx2,
+                                            torch::Tensor grad_dist1,
+                                            torch::Tensor grad_dist2) {
+  return chamfer_cuda_backward(xyz1, xyz2, idx1, idx2, grad_dist1, grad_dist2);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &chamfer_forward, "Chamfer forward (CUDA)");
+  m.def("backward", &chamfer_backward, "Chamfer backward (CUDA)");
+}

openpoints-0.1.0/cpp/chamfer_dist/setup.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+# @Author: Haozhe Xie
+# @Date:   2019-08-07 20:54:24
+# @Last Modified by:   Haozhe Xie
+# @Last Modified time: 2019-12-10 10:04:25
+# @Email:  cshzxie@gmail.com
+
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(name='chamfer',
+      version='2.0.0',
+      ext_modules=[
+          CUDAExtension('chamfer', [
+              'chamfer_cuda.cpp',
+              'chamfer.cu',
+          ]),
+      ],
+      cmdclass={'build_ext': BuildExtension})