Skip to content

GitLab

Commit 1111319d authored by Alexander Liao's avatar Alexander Liao
Browse files

options support; correctness across dimensions; more testing

parent 25d62f0a
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 682 additions and 49 deletions
csrc/cpu/radius_cpu.cpp
View file @ 1111319d
#include "radius_cpu.h"
#include <algorithm>
#include "utils.h"
torch::Tensor radius_cpu(torch::Tensor q, torch::Tensor s,
torch::Tensor ptr_x, torch::Tensor ptr_y,
torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
float radius, int max_num){
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().reshape({-1, 1});
auto batch_y = ptr_y.clone().reshape({-1, 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);
CHECK_CPU(query);
CHECK_CPU(support);
torch::Tensor out;
std::vector<long> neighbors_indices;
......@@ -58,6 +43,7 @@ torch::Tensor radius_cpu(torch::Tensor q, torch::Tensor s,
return result;
}
void get_size_batch(const vector<long>& batch, vector<long>& res){
res.resize(batch[batch.size()-1]-batch[0]+1, 0);
......@@ -74,4 +60,54 @@ void get_size_batch(const vector<long>& batch, vector<long>& res){
}
}
res[ind-batch[0]] = incr;
}
torch::Tensor batch_radius_cpu(torch::Tensor query,
torch::Tensor support,
torch::Tensor query_batch,
torch::Tensor support_batch,
float radius, int max_num) {
torch::Tensor out;
auto data_qb = query_batch.data_ptr<long>();
auto data_sb = support_batch.data_ptr<long>();
std::vector<long> query_batch_stl = std::vector<long>(data_qb, data_qb+query_batch.size(0));
std::vector<long> size_query_batch_stl;
get_size_batch(query_batch_stl, size_query_batch_stl);
std::vector<long> support_batch_stl = std::vector<long>(data_sb, data_sb+support_batch.size(0));
std::vector<long> size_support_batch_stl;
get_size_batch(support_batch_stl, size_support_batch_stl);
std::vector<long> neighbors_indices;
auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
int max_count = 0;
AT_DISPATCH_ALL_TYPES(query.scalar_type(), "batch_radius_search", [&] {
auto data_q = query.data_ptr<scalar_t>();
auto data_s = support.data_ptr<scalar_t>();
std::vector<scalar_t> queries_stl = std::vector<scalar_t>(data_q,
data_q + query.size(0)*query.size(1));
std::vector<scalar_t> supports_stl = std::vector<scalar_t>(data_s,
data_s + support.size(0)*support.size(1));
int dim = torch::size(query, 1);
max_count = batch_nanoflann_neighbors<scalar_t>(queries_stl,
supports_stl,
size_query_batch_stl,
size_support_batch_stl,
neighbors_indices,
radius,
dim,
max_num
);
});
long* neighbors_indices_ptr = neighbors_indices.data();
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();
return out.clone();
}
\ No newline at end of file
csrc/cpu/radius_cpu.h
View file @ 1111319d
......@@ -7,5 +7,10 @@
#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);
\ No newline at end of file
float radius, int max_num);
torch::Tensor batch_radius_cpu(torch::Tensor query,
torch::Tensor support,
torch::Tensor query_batch,
torch::Tensor support_batch,
float radius, int max_num);
\ No newline at end of file
csrc/cpu/utils/cloud.h
View file @ 1111319d
......@@ -24,18 +24,13 @@ struct PointCloud
void set(std::vector<scalar_t> new_pts, int dim){
// 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(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 (size_t j = 0; j < (size_t)dim; j++) {
point[j]=new_pts[i+j];
//point.pt[j] = new_pts[i+j];
}
temp[i/dim] = point;
}
......@@ -46,7 +41,6 @@ struct PointCloud
void set_batch(std::vector<scalar_t> new_pts, int begin, int size, int dim){
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);
std::vector<scalar_t> point(dim);
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 @ 1111319d
......@@ -40,10 +40,10 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
// Search params
nanoflann::SearchParams search_params;
search_params.sorted = true;
// search_params.sorted = true;
std::vector< std::vector<std::pair<size_t, scalar_t> > > list_matches(pcd_query.pts.size());
float eps = 0.00001;
float eps = 0.000001;
// indices
size_t i0 = 0;
......@@ -61,6 +61,12 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
std::vector<std::pair<size_t, scalar_t> > ret_matches;
const size_t nMatches = index->radiusSearch(query_pt, search_radius+eps, ret_matches, search_params);
//cout << "radiusSearch(): radius=" << search_radius << " -> " << nMatches << " matches\n";
//for (size_t i = 0; i < nMatches; i++)
// cout << "idx["<< i << "]=" << ret_matches[i].first << " dist["<< i << "]=" << ret_matches[i].second << endl;
//cout << "\n";
list_matches[i0] = ret_matches;
if(max_count < nMatches) max_count = nMatches;
i0++;
......@@ -107,4 +113,139 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
}
\ No newline at end of file
}
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){
// Initiate variables
// ******************
// indices
int i0 = 0;
// Square radius
const scalar_t r2 = static_cast<scalar_t>(radius*radius);
// Counting vector
int max_count = 0;
float d2;
// batch index
long b = 0;
long sum_qb = 0;
long sum_sb = 0;
float eps = 0.000001;
// 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+eps, 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
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;
}
csrc/cpu/utils/neighbors.h
View file @ 1111319d
......@@ -10,4 +10,12 @@ 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);
\ No newline at end of file
vector<long>& neighbors_indices, float radius, int dim, int max_num);
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);
\ No newline at end of file
csrc/radius.cpp
View file @ 1111319d
......@@ -10,16 +10,48 @@
PyMODINIT_FUNC PyInit__radius(void) { return NULL; }
#endif
torch::Tensor radius(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
torch::Tensor ptr_y, double r, int64_t max_num_neighbors) {
torch::Tensor radius(torch::Tensor x, torch::Tensor y, torch::optional<torch::Tensor> ptr_x,
torch::optional<torch::Tensor> ptr_y, double r, int64_t max_num_neighbors) {
if (x.device().is_cuda()) {
#ifdef WITH_CUDA
return radius_cuda(x, y, ptr_x, ptr_y, r, max_num_neighbors);
if (!(ptr_x.has_value()) && !(ptr_y.has_value())) {
auto batch_x = torch::tensor({0,torch::size(x,0)}).to(torch::kLong).to(torch::kCUDA);
auto batch_y = torch::tensor({0,torch::size(y,0)}).to(torch::kLong).to(torch::kCUDA);
return radius_cuda(x, y, batch_x, batch_y, r, max_num_neighbors);
}
else if (!(ptr_x.has_value())) {
auto batch_x = torch::tensor({0,torch::size(x,0)}).to(torch::kLong).to(torch::kCUDA);
auto batch_y = ptr_y.value();
return radius_cuda(x, y, batch_x, batch_y, r, max_num_neighbors);
}
else if (!(ptr_y.has_value())) {
auto batch_x = ptr_x.value();
auto batch_y = torch::tensor({0,torch::size(y,0)}).to(torch::kLong).to(torch::kCUDA);
return radius_cuda(x, y, batch_x, batch_y, r, max_num_neighbors);
}
auto batch_x = ptr_x.value();
auto batch_y = ptr_y.value();
return radius_cuda(x, y, batch_x, batch_y, r, max_num_neighbors);
#else
AT_ERROR("Not compiled with CUDA support");
#endif
} else {
return radius_cpu(x, y, ptr_x, ptr_y, r, max_num_neighbors);
if (!(ptr_x.has_value()) && !(ptr_y.has_value())) {
return radius_cpu(x,y,r,max_num_neighbors);
}
if (!(ptr_x.has_value())) {
auto batch_x = torch::zeros({torch::size(x,0)}).to(torch::kLong);
auto batch_y = ptr_y.value();
return batch_radius_cpu(x, y, batch_x, batch_y, r, max_num_neighbors);
}
else if (!(ptr_y.has_value())) {
auto batch_x = ptr_x.value();
auto batch_y = torch::zeros({torch::size(y,0)}).to(torch::kLong);
return batch_radius_cpu(x, y, batch_x, batch_y, r, max_num_neighbors);
}
auto batch_x = ptr_x.value();
auto batch_y = ptr_y.value();
return batch_radius_cpu(x, y, batch_x, batch_y, r, max_num_neighbors);
}
}
......
test/test_radius.py
View file @ 1111319d
This diff is collapsed.
torch_cluster/radius.py
View file @ 1111319d
from typing import Optional
import torch
import scipy
def sample(col, count):
if col.size(0) > count:
col = col[torch.randperm(col.size(0))][:count]
return col
def radius(x: torch.Tensor, y: torch.Tensor, r: float,
batch_x: Optional[torch.Tensor] = None,
......@@ -50,7 +55,7 @@ def radius(x: torch.Tensor, y: torch.Tensor, r: float,
ptr_x = deg.new_zeros(batch_size + 1)
torch.cumsum(deg, 0, out=ptr_x[1:])
else:
ptr_x = torch.tensor([0, x.size(0)], device=x.device)
ptr_x = None#torch.tensor([0, x.size(0)], device=x.device)
if batch_y is not None:
assert y.size(0) == batch_y.numel()
......@@ -61,28 +66,33 @@ def radius(x: torch.Tensor, y: torch.Tensor, r: float,
ptr_y = deg.new_zeros(batch_size + 1)
torch.cumsum(deg, 0, out=ptr_y[1:])
else:
ptr_y = torch.tensor([0, y.size(0)], device=y.device)
ptr_y = None#torch.tensor([0, y.size(0)], device=y.device)
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)
#if batch_x is None:
# batch_x = x.new_zeros(x.size(0), dtype=torch.long)
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
#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)
#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.dim() == 2
if batch_x is not None:
assert batch_x.dim() == 1
assert x.size(0) == batch_x.size(0)
assert y.dim() == 2
if batch_y is not None:
assert batch_y.dim() == 1
assert y.size(0) == batch_y.size(0)
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
result = torch.ops.torch_cluster.radius(x, y, batch_x, batch_y, r,
max_num_neighbors)
max_num_neighbors)
return result
......
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