Merge pull request #22 from nicolas-chaulet/cpuversion

Cpuversion

Author: tchaton

.devcontainer/Dockerfile

@@ -0,0 +1,33 @@
+#-------------------------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See https://go.microsoft.com/fwlink/?linkid=2090316 for license information.
+#-------------------------------------------------------------------------------------------------------------
+
+FROM ubuntu:bionic
+
+# Avoid warnings by switching to noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This Dockerfile adds a non-root user with sudo access. Use the "remoteUser"
+# property in devcontainer.json to use it. On Linux, the container user's GID/UIDs
+# will be updated to match your local UID/GID (when using the dockerFile property).
+# See https://aka.ms/vscode-remote/containers/non-root-user for details.
+ARG USERNAME=vscode
+ARG USER_UID=1000
+ARG USER_GID=$USER_UID
+
+# Uncomment the following COPY line and the corresponding lines in the `RUN` command if you wish to
+# include your requirements in the image itself. It is suggested that you only do this if your
+# requirements rarely (if ever) change.
+
+RUN apt-get update \
+    && apt-get install -y --fix-missing --no-install-recommends\
+    libffi-dev libssl-dev build-essential \
+    python3-pip python3-dev python3-venv python3-setuptools\
+    git iproute2 procps lsb-release clang-format \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN pip3 install -U pip
+RUN pip3 install torch numpy scikit-learn flake8 setuptools
+RUN pip3 install torch_cluster torch_sparse torch_scatter torch_geometric

.devcontainer/devcontainer.json

@@ -0,0 +1,26 @@
+// For format details, see https://aka.ms/vscode-remote/devcontainer.json or this file's README at:
+// https://github.com/microsoft/vscode-dev-containers/tree/v0.101.1/containers/python-3
+{
+	"name": "Python 3",
+	"context": "..",
+	"dockerFile": "Dockerfile",
+	// Set *default* container specific settings.json values on container create.
+	"settings": {
+		"terminal.integrated.shell.linux": "/bin/bash",
+		"python.pythonPath": "/usr/local/bin/python",
+		"python.linting.enabled": true,
+		"python.linting.pylintEnabled": true,
+		"python.linting.pylintPath": "/usr/local/bin/pylint"
+	},
+	// Add the IDs of extensions you want installed when the container is created.
+	"extensions": [
+		"ms-python.python",
+		"ms-vscode.cpptools"
+	]
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+	// Use 'postCreateCommand' to run commands after the container is created.
+	// "postCreateCommand": "pip install -r requirements.txt",
+	// Uncomment to connect as a non-root user. See https://aka.ms/vscode-remote/containers/non-root.
+	// "remoteUser": "vscode"
+}

cpu/include/fps.h

@@ -0,0 +1,3 @@
+#pragma once
+#include <torch/extension.h>
+at::Tensor fps(at::Tensor points, const int nsamples, bool random = true);

cpu/include/group_points.h

@@ -0,0 +1,7 @@
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor knn_interpolate(at::Tensor features, at::Tensor idx, at::Tensor weight);
+
+at::Tensor knn_interpolate_grad(at::Tensor grad_out, at::Tensor idx, at::Tensor weight,
+                                const int m);

cpu/include/interpolate.h

@@ -0,0 +1,3 @@
+#pragma once
+#include <torch/extension.h>
+std::pair<at::Tensor, at::Tensor> dense_knn(at::Tensor support, at::Tensor query, int k);

cpu/include/knn.h

@@ -17,3 +17,7 @@ int batch_nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& suppo
                               vector<long>& q_batches, vector<long>& s_batches,
                               vector<long>& neighbors_indices, vector<float>& dists, float radius,
                               int max_num, int mode);
+
+template <typename scalar_t>
+void nanoflann_knn_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
+                             vector<long>& neighbors_indices, vector<float>& dists, int k);

cpu/include/neighbors.h

@@ -3,4 +3,4 @@
 
 #define CHECK_CPU(x) AT_ASSERTM(!x.type().is_cuda(), #x " must be a CPU tensor")
 
-#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be a contiguous tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be a contiguous tensor")

cpu/include/utils.h

@@ -1,11 +1,17 @@
 #include "ball_query.h"
-#include "group_points.h"
+#include "fps.h"
+#include "interpolate.h"
+#include "knn.h"
 
 using namespace pybind11::literals;
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 {
-    m.def("group_points", &group_points);
+    m.def("dense_knn", &dense_knn, "", "support"_a, "querry"_a, "k"_a);
+    m.def("knn_interpolate", &knn_interpolate, "", "features"_a, "idx"_a, "weights"_a);
+    m.def("knn_interpolate_grad", &knn_interpolate_grad, "", "grad_out"_a, "idx"_a, "weights"_a,
+          "m"_a);
+    m.def("fps", &fps, "", "points"_a, "num_samples"_a, "random"_a);
 
     m.def("ball_query", &ball_query,
           "compute the radius search of a point cloud using nanoflann"