Skip to content

GitLab

Commit 59dfc410 authored by Alexander Liao's avatar Alexander Liao
Browse files

cleaning up according to flake8; c++ api now working (except options)

parent 303e889a
Hide whitespace changes
Inline Side-by-side
Showing with 50 additions and 264 deletions
csrc/cpu/radius_cpu.cpp
View file @ 59dfc410
......@@ -2,25 +2,26 @@
#include "utils.h"
torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
torch::Tensor radius_cpu(torch::Tensor q, torch::Tensor s,
torch::Tensor ptr_x, torch::Tensor ptr_y,
float radius, int max_num){
CHECK_CPU(query);
CHECK_CPU(support);
CHECK_CPU(q);
CHECK_CPU(s);
/*
x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
auto batch_x = ptr_x.clone();
auto batch_y = ptr_y.clone();
*/
batch_x._mul(2*radius);
batch_y._mul(2*radius);
auto batch_x = ptr_x.clone().reshape({-1, 1});
auto batch_y = ptr_y.clone().reshape({-1, 1});
auto query = torch::cat({query,batch_x},-1);
auto support = torch::cat({support,batch_y},-1);
*/
batch_x.mul_(2*radius);
batch_y.mul_(2*radius);
auto query = torch::cat({q,batch_x},-1);
auto support = torch::cat({s,batch_y},-1);
torch::Tensor out;
std::vector<long> neighbors_indices;
......@@ -43,9 +44,18 @@ torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
});
long* neighbors_indices_ptr = neighbors_indices.data();
out = torch::from_blob(neighbors_indices_ptr, {neighbors_indices.size()/2, 2}, options=options);
return out.t().clone();
const long long tsize = static_cast<long long>(neighbors_indices.size()/2);
out = torch::from_blob(neighbors_indices_ptr, {tsize, 2}, options=options);
out = out.t();
auto result = torch::zeros_like(out);
auto index = torch::tensor({0,1});
result.index_copy_(0, index, out);
return result;
}
void get_size_batch(const vector<long>& batch, vector<long>& res){
......
csrc/cpu/radius_cpu.h
View file @ 59dfc410
#pragma once
#include <torch/extension.h>
#include <ATen/ATen.h>
#include "utils/neighbors.h"
#include "utils/neighbors.cpp"
#include <iostream>
#include "compat.h"
torch::Tensor radius_cpu(torch::Tensor query,
torch::Tensor support,torch::Tensor ptr_x,
torch::Tensor ptr_y,
float radius, int max_num);
/*
using namespace pybind11::literals;
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("radius_search",
&radius_search,
"compute the radius search of a point cloud using nanoflann"
"-query : a pytorch tensor of size N1 x d,. used to query the nearest neighbors"
"- support : a pytorch tensor of size N2 x d. used to build the tree"
"- radius : float number, size of the ball for the radius search."
"- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors will be computed). "
" - mode : int number that indicate which format for the neighborhood"
" mode=0 mean a matrix of neighbors(-1 for shadow neighbors)"
"mode=1 means a matrix of edges of size Num_edge x 2"
"return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2.",
"query"_a, "support"_a, "radius"_a, "dim"_a, "max_num"_a=-1, "mode"_a=0);
m.def("batch_radius_search",
&batch_radius_search,
"compute the radius search of a point cloud for each batch using nanoflann"
"-query : a pytorch tensor (float) of size N1 x d,. used to query the nearest neighbors"
"- support : a pytorch tensor(float) of size N2 x d. used to build the tree"
"- query_batch : a pytorch tensor(long) contains indices of the batch of the query size N1"
"NB : the batch must be sorted"
"- support_batch: a pytorch tensor(long) contains indices of the batch of the support size N2"
"NB: the batch must be sorted"
"-radius: float number, size of the ball for the radius search."
"- max_num : int number, indicate the maximum of neaghbors allowed(if -1 then all the possible neighbors wrt the radius will be computed)."
"- mode : int number that indicate which format for the neighborhood"
"mode=0 mean a matrix of neighbors(N2 for shadow neighbors)"
"mode=1 means a matrix of edges of size Num_edge x 2"
"return a tensor of size N1 x M where M is either max_num or the maximum number of neighbors found if mode = 0, if mode=1 return a tensor of size Num_edge x 2.",
"query"_a, "support"_a, "query_batch"_a, "support_batch"_a, "radius"_a, "dim"_a, "max_num"_a=-1, "mode"_a=0);
}
*/
\ No newline at end of file
torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
torch::Tensor ptr_x, torch::Tensor ptr_y,
float radius, int max_num);
\ No newline at end of file
csrc/cpu/utils/cloud.h
View file @ 59dfc410
......@@ -26,14 +26,14 @@ struct PointCloud
// pts = std::vector<Point>((Point*)new_pts, (Point*)new_pts+new_pts.size()/3);
std::vector<std::vector<scalar_t>> temp(new_pts.size()/dim);
for(unsigned int i=0; i < new_pts.size(); i++){
for(size_t i=0; i < new_pts.size(); i++){
if(i%dim == 0){
//Point point;
std::vector<scalar_t> point(dim);
//std::vector<scalar_t> vect(sizeof(scalar_t)*dim, 0)
//point.pt = temp;
for (unsigned int j = 0; j < dim; j++) {
for (size_t j = 0; j < (size_t)dim; j++) {
point[j]=new_pts[i+j];
//point.pt[j] = new_pts[i+j];
}
......@@ -47,9 +47,8 @@ struct PointCloud
std::vector<std::vector<scalar_t>> temp(size);
for(int i=0; i < size; i++){
//std::vector<scalar_t> temp(sizeof(scalar_t)*dim, 0);
//point.pt = temp;
std::vector<scalar_t> point(dim);
for (unsigned int j = 0; j < dim; j++) {
for (size_t j = 0; j < (size_t)dim; j++) {
point[j] = new_pts[dim*(begin+i)+j];
}
......
csrc/cpu/utils/neighbors.cpp
View file @ 59dfc410
// Taken from https://github.com/HuguesTHOMAS/KPConv
// 3D Version https://github.com/HuguesTHOMAS/KPConv
#include "neighbors.h"
......@@ -10,16 +10,10 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
// Initiate variables
// ******************
// square radius
const scalar_t search_radius = static_cast<scalar_t>(radius*radius);
// indices
int i0 = 0;
// Counting vector
int max_count = 1;
float d2;
size_t max_count = 1;
// Nanoflann related variables
// ***************************
......@@ -32,7 +26,7 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
pcd_query.set(queries, dim);
// Tree parameters
nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
nanoflann::KDTreeSingleIndexAdaptorParams tree_params(15 /* max leaf */);
// KDTree type definition
typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
......@@ -51,6 +45,9 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
float eps = 0.00001;
// indices
size_t i0 = 0;
for (auto& p0 : pcd_query.pts){
// Find neighbors
......@@ -62,7 +59,6 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
list_matches[i0].reserve(max_count);
std::vector<std::pair<size_t, scalar_t> > ret_matches;
//nanoflann::RadiusResultSet<float,size_t> resultSet(r2, indices_dists);
const size_t nMatches = index->radiusSearch(&query_pt[0], search_radius+eps, ret_matches, search_params);
list_matches[i0] = ret_matches;
......@@ -84,19 +80,19 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
max_count = max_num;
}
int size = 0; // total number of edges
size_t size = 0; // total number of edges
for (auto& inds : list_matches){
if(inds.size() <= max_count)
size += inds.size();
else
size += max_count;
}
neighbors_indices.resize(size*2);
int i1 = 0; // index of the query points
int u = 0; // curent index of the neighbors_indices
size_t i1 = 0; // index of the query points
size_t u = 0; // curent index of the neighbors_indices
for (auto& inds : list_matches){
for (int j = 0; j < max_count; j++){
for (size_t j = 0; j < max_count; j++){
if(j < inds.size()){
neighbors_indices[u] = inds[j].first;
neighbors_indices[u + 1] = i1;
......@@ -111,165 +107,4 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
}
template<typename scalar_t>
int batch_nanoflann_neighbors (vector<scalar_t>& queries,
vector<scalar_t>& supports,
vector<long>& q_batches,
vector<long>& s_batches,
vector<long>& neighbors_indices,
float radius, int dim, int max_num,
int mode){
// Initiate variables
// ******************
// indices
int i0 = 0;
// Square radius
float r2 = radius * radius;
// Counting vector
int max_count = 0;
float d2;
// batch index
long b = 0;
long sum_qb = 0;
long sum_sb = 0;
// Nanoflann related variables
// ***************************
// CLoud variable
PointCloud<scalar_t> current_cloud;
PointCloud<scalar_t> query_pcd;
query_pcd.set(queries, dim);
vector<vector<pair<size_t, scalar_t> > > all_inds_dists(query_pcd.pts.size());
// Tree parameters
nanoflann::KDTreeSingleIndexAdaptorParams tree_params(10 /* max leaf */);
// KDTree type definition
typedef nanoflann::KDTreeSingleIndexAdaptor< nanoflann::L2_Adaptor<scalar_t, PointCloud<scalar_t> > , PointCloud<scalar_t>> my_kd_tree_t;
// Pointer to trees
my_kd_tree_t* index;
// Build KDTree for the first batch element
current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
index = new my_kd_tree_t(dim, current_cloud, tree_params);
index->buildIndex();
// Search neigbors indices
// ***********************
// Search params
nanoflann::SearchParams search_params;
search_params.sorted = true;
for (auto& p0 : query_pcd.pts){
// Check if we changed batch
scalar_t query_pt[dim];
std::copy(p0.begin(), p0.end(), query_pt);
/*
std::cout << "\n ========== \n";
for(int i=0; i < dim; i++)
std::cout << query_pt[i] << '\n';
std::cout << "\n ========== \n";
*/
if (i0 == sum_qb + q_batches[b]){
sum_qb += q_batches[b];
sum_sb += s_batches[b];
b++;
// Change the points
current_cloud.pts.clear();
current_cloud.set_batch(supports, sum_sb, s_batches[b], dim);
// Build KDTree of the current element of the batch
delete index;
index = new my_kd_tree_t(dim, current_cloud, tree_params);
index->buildIndex();
}
// Initial guess of neighbors size
all_inds_dists[i0].reserve(max_count);
// Find neighbors
size_t nMatches = index->radiusSearch(query_pt, r2, all_inds_dists[i0], search_params);
// Update max count
std::vector<std::pair<size_t, float> > indices_dists;
nanoflann::RadiusResultSet<float,size_t> resultSet(r2, indices_dists);
index->findNeighbors(resultSet, query_pt, search_params);
if (nMatches > max_count)
max_count = nMatches;
// Increment query idx
i0++;
}
// how many neighbors do we keep
if(max_num > 0) {
max_count = max_num;
}
// Reserve the memory
if(mode == 0){
neighbors_indices.resize(query_pcd.pts.size() * max_count);
i0 = 0;
sum_sb = 0;
sum_qb = 0;
b = 0;
for (auto& inds_dists : all_inds_dists){// Check if we changed batch
if (i0 == sum_qb + q_batches[b]){
sum_qb += q_batches[b];
sum_sb += s_batches[b];
b++;
}
for (int j = 0; j < max_count; j++){
if (j < inds_dists.size())
neighbors_indices[i0 * max_count + j] = inds_dists[j].first + sum_sb;
else
neighbors_indices[i0 * max_count + j] = supports.size();
}
i0++;
}
delete index;
}
else if(mode == 1){
int size = 0; // total number of edges
for (auto& inds_dists : all_inds_dists){
if(inds_dists.size() <= max_count)
size += inds_dists.size();
else
size += max_count;
}
neighbors_indices.resize(size * 2);
i0 = 0;
sum_sb = 0;
sum_qb = 0;
b = 0;
int u = 0;
for (auto& inds_dists : all_inds_dists){
if (i0 == sum_qb + q_batches[b]){
sum_qb += q_batches[b];
sum_sb += s_batches[b];
b++;
}
for (int j = 0; j < max_count; j++){
if (j < inds_dists.size()){
neighbors_indices[u] = inds_dists[j].first + sum_sb;
neighbors_indices[u + 1] = i0;
u += 2;
}
}
i0++;
}
}
return max_count;
}
}
\ No newline at end of file
csrc/cpu/utils/neighbors.h
View file @ 59dfc410
......@@ -10,13 +10,4 @@ using namespace std;
template<typename scalar_t>
int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
vector<long>& neighbors_indices, float radius, int dim, int max_num, int mode);
template<typename scalar_t>
int batch_nanoflann_neighbors(vector<scalar_t>& queries,
vector<scalar_t>& supports,
vector<long>& q_batches,
vector<long>& s_batches,
vector<long>& neighbors_indices,
float radius, int dim, int max_num, int mode);
vector<long>& neighbors_indices, float radius, int dim, int max_num, int mode);
\ No newline at end of file
torch_cluster/radius.py
View file @ 59dfc410
from typing import Optional
import torch
@torch.jit.script
def sample(col: torch.Tensor, count: int) -> torch.Tensor:
if col.size(0) > count:
col = col[torch.randperm(col.size(0), dtype=torch.long)][:count]
return col
def radius(x: torch.Tensor, y: torch.Tensor, r: float,
batch_x: Optional[torch.Tensor] = None,
......@@ -70,8 +63,8 @@ def radius(x: torch.Tensor, y: torch.Tensor, r: float,
else:
ptr_y = torch.tensor([0, y.size(0)], device=y.device)
return torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
max_num_neighbors)
result = torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
max_num_neighbors)
else:
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
......@@ -79,27 +72,19 @@ def radius(x: torch.Tensor, y: torch.Tensor, r: float,
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
batch_x = batch_x.to(x.dtype)
batch_y = batch_y.to(y.dtype)
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
result = torch.ops.torch_cluster.radius(x, y, batch_x, batch_y, r,
max_num_neighbors)
return torch.ops.torch_cluster.radius(x, y, x, y, r,
max_num_neighbors)
"""
tree = scipy.spatial.cKDTree(x.detach().numpy())
col = tree.query_ball_point(y.detach().numpy(), r)
col = [torch.tensor(c, dtype=torch.long) for c in col]
col = [sample(c, max_num_neighbors) for c in col]
row = [torch.full_like(c, i) for i, c in enumerate(col)]
row, col = torch.cat(row, dim=0), torch.cat(col, dim=0)
mask = col < int(tree.n)
return torch.stack([row[mask], col[mask]], dim=0)
"""
return result
def radius_graph(x: torch.Tensor, r: float,
......@@ -138,6 +123,7 @@ def radius_graph(x: torch.Tensor, r: float,
assert flow in ['source_to_target', 'target_to_source']
row, col = radius(x, x, r, batch, batch,
max_num_neighbors if loop else max_num_neighbors + 1)
if x.is_cuda:
row, col = (col, row) if flow == 'source_to_target' else (row, col)
else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment