Skip to content

Commit 962fc02

Browse files
liaopeiyuanliaopeiyuan
committed
attempting to fix windows build memory error
1 parent d4fe021 commit 962fc02

File tree

6 files changed

+48
-46
lines changed

6 files changed

+48
-46
lines changed

‎csrc/cpu/radius_cpu.cpp‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
11
#include "radius_cpu.h"
22
#include <algorithm>
33
#include "utils.h"
4+
#include <cstdint>
5+
46

57
torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
6-
float radius, int max_num){
8+
double radius, int64_t max_num){
79

810
CHECK_CPU(query);
911
CHECK_CPU(support);
1012

1113
torch::Tensor out;
12-
std::vector<long> neighbors_indices;
14+
std::vector<size_t>* neighbors_indices = new std::vector<size_t>();
1315
auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
1416
int max_count = 0;
1517

@@ -28,9 +30,9 @@ torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
2830

2931
});
3032

31-
long* neighbors_indices_ptr = neighbors_indices.data();
33+
size_t* neighbors_indices_ptr = neighbors_indices->data();
3234

33-
const long long tsize = static_cast<long long>(neighbors_indices.size()/2);
35+
const long long tsize = static_cast<long long>(neighbors_indices->size()/2);
3436
out = torch::from_blob(neighbors_indices_ptr, {tsize, 2}, options=options);
3537
out = out.t();
3638

@@ -60,7 +62,7 @@ torch::Tensor batch_radius_cpu(torch::Tensor query,
6062
torch::Tensor support,
6163
torch::Tensor query_batch,
6264
torch::Tensor support_batch,
63-
float radius, int max_num) {
65+
double radius, int64_t max_num) {
6466

6567
torch::Tensor out;
6668
auto data_qb = query_batch.data_ptr<int64_t>();
@@ -71,7 +73,7 @@ torch::Tensor batch_radius_cpu(torch::Tensor query,
7173
std::vector<long> support_batch_stl = std::vector<long>(data_sb, data_sb+support_batch.size(0));
7274
std::vector<long> size_support_batch_stl;
7375
get_size_batch(support_batch_stl, size_support_batch_stl);
74-
std::vector<long> neighbors_indices;
76+
std::vector<size_t>* neighbors_indices = new std::vector<size_t>();
7577
auto options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
7678
int max_count = 0;
7779

@@ -95,10 +97,10 @@ torch::Tensor batch_radius_cpu(torch::Tensor query,
9597
);
9698
});
9799

98-
long* neighbors_indices_ptr = neighbors_indices.data();
100+
size_t* neighbors_indices_ptr = neighbors_indices->data();
99101

100102

101-
const long long tsize = static_cast<long long>(neighbors_indices.size()/2);
103+
const long long tsize = static_cast<long long>(neighbors_indices->size()/2);
102104
out = torch::from_blob(neighbors_indices_ptr, {tsize, 2}, options=options);
103105
out = out.t();
104106

‎csrc/cpu/radius_cpu.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@
77
#include "compat.h"
88

99
torch::Tensor radius_cpu(torch::Tensor query, torch::Tensor support,
10-
float radius, int max_num);
10+
double radius, int64_t max_num);
1111

1212
torch::Tensor batch_radius_cpu(torch::Tensor query,
1313
torch::Tensor support,
1414
torch::Tensor query_batch,
1515
torch::Tensor support_batch,
16-
float radius, int max_num);
16+
double radius, int64_t max_num);

‎csrc/cpu/utils/cloud.h‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,30 +20,30 @@
2020
template<typename scalar_t>
2121
struct PointCloud
2222
{
23-
std::vector<std::vector<scalar_t>> pts;
23+
std::vector<std::vector<scalar_t>*> pts;
2424

2525
void set(std::vector<scalar_t> new_pts, int dim){
2626

27-
std::vector<std::vector<scalar_t>> temp(new_pts.size()/dim);
27+
std::vector<std::vector<scalar_t>*> temp(new_pts.size()/dim);
2828
for(size_t i=0; i < new_pts.size(); i++){
2929
if(i%dim == 0){
30-
std::vector<scalar_t> point(dim);
30+
std::vector<scalar_t>* point = new std::vector<scalar_t>(dim);
3131

3232
for (size_t j = 0; j < (size_t)dim; j++) {
33-
point[j]=new_pts[i+j];
33+
(*point)[j]=new_pts[i+j];
3434
}
3535
temp[i/dim] = point;
3636
}
3737
}
3838

3939
pts = temp;
4040
}
41-
void set_batch(std::vector<scalar_t> new_pts, int begin, int size, int dim){
42-
std::vector<std::vector<scalar_t>> temp(size);
43-
for(int i=0; i < size; i++){
44-
std::vector<scalar_t> point(dim);
41+
void set_batch(std::vector<scalar_t> new_pts, size_t begin, long size, int dim){
42+
std::vector<std::vector<scalar_t>*> temp(size);
43+
for(size_t i=0; i < (size_t)size; i++){
44+
std::vector<scalar_t>* point = new std::vector<scalar_t>(dim);
4545
for (size_t j = 0; j < (size_t)dim; j++) {
46-
point[j] = new_pts[dim*(begin+i)+j];
46+
(*point)[j] = new_pts[dim*(begin+i)+j];
4747
}
4848

4949
temp[i] = point;
@@ -58,7 +58,7 @@ struct PointCloud
5858
// Returns the dim'th component of the idx'th point in the class:
5959
inline scalar_t kdtree_get_pt(const size_t idx, const size_t dim) const
6060
{
61-
return pts[idx][dim];
61+
return (*pts[idx])[dim];
6262
}
6363

6464
// Optional bounding-box computation: return false to default to a standard bbox computation loop.

‎csrc/cpu/utils/neighbors.cpp‎

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
#include "neighbors.h"
55

66
template<typename scalar_t>
7-
int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
8-
vector<long>& neighbors_indices, float radius, int dim, int max_num){
7+
size_t nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
8+
vector<size_t>*& neighbors_indices, double radius, int dim, int64_t max_num){
99

1010
const scalar_t search_radius = static_cast<scalar_t>(radius*radius);
1111

@@ -37,21 +37,21 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
3737
// search_params.sorted = true;
3838
std::vector< std::vector<std::pair<size_t, scalar_t> > > list_matches(pcd_query.pts.size());
3939

40-
float eps = 0.000001;
40+
double eps = 0.000001;
4141

4242
// indices
4343
size_t i0 = 0;
4444

45-
for (auto& p0 : pcd_query.pts){
46-
45+
for (auto& p : pcd_query.pts){
46+
auto p0 = *p;
4747
// Find neighbors
4848
scalar_t* query_pt = new scalar_t[dim];
4949
std::copy(p0.begin(), p0.end(), query_pt);
5050

5151
list_matches[i0].reserve(max_count);
5252
std::vector<std::pair<size_t, scalar_t> > ret_matches;
5353

54-
const size_t nMatches = index->radiusSearch(query_pt, search_radius+eps, ret_matches, search_params);
54+
const size_t nMatches = index->radiusSearch(query_pt, (scalar_t)(search_radius+eps), ret_matches, search_params);
5555

5656
list_matches[i0] = ret_matches;
5757
if(max_count < nMatches) max_count = nMatches;
@@ -71,14 +71,14 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
7171
size += max_count;
7272
}
7373

74-
neighbors_indices.resize(size*2);
74+
neighbors_indices->resize(size*2);
7575
size_t i1 = 0; // index of the query points
7676
size_t u = 0; // curent index of the neighbors_indices
7777
for (auto& inds : list_matches){
7878
for (size_t j = 0; j < max_count; j++){
7979
if(j < inds.size()){
80-
neighbors_indices[u] = inds[j].first;
81-
neighbors_indices[u + 1] = i1;
80+
(*neighbors_indices)[u] = inds[j].first;
81+
(*neighbors_indices)[u + 1] = i1;
8282
u += 2;
8383
}
8484
}
@@ -93,12 +93,12 @@ int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
9393
}
9494

9595
template<typename scalar_t>
96-
int batch_nanoflann_neighbors (vector<scalar_t>& queries,
96+
size_t batch_nanoflann_neighbors (vector<scalar_t>& queries,
9797
vector<scalar_t>& supports,
9898
vector<long>& q_batches,
9999
vector<long>& s_batches,
100-
vector<long>& neighbors_indices,
101-
float radius, int dim, int max_num){
100+
vector<size_t>*& neighbors_indices,
101+
double radius, int dim, int64_t max_num){
102102

103103

104104
// Initiate variables
@@ -117,7 +117,7 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
117117
size_t sum_qb = 0;
118118
size_t sum_sb = 0;
119119

120-
float eps = 0.000001;
120+
double eps = 0.000001;
121121
// Nanoflann related variables
122122
// ***************************
123123

@@ -145,7 +145,8 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
145145
nanoflann::SearchParams search_params;
146146
search_params.sorted = true;
147147

148-
for (auto& p0 : query_pcd.pts){
148+
for (auto& p : query_pcd.pts){
149+
auto p0 = *p;
149150
// Check if we changed batch
150151

151152
scalar_t* query_pt = new scalar_t[dim];
@@ -193,7 +194,7 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
193194
else
194195
size += max_count;
195196
}
196-
neighbors_indices.resize(size * 2);
197+
neighbors_indices->resize(size * 2);
197198
i0 = 0;
198199
sum_sb = 0;
199200
sum_qb = 0;
@@ -207,8 +208,8 @@ int batch_nanoflann_neighbors (vector<scalar_t>& queries,
207208
}
208209
for (size_t j = 0; j < max_count; j++){
209210
if (j < inds_dists.size()){
210-
neighbors_indices[u] = inds_dists[j].first + sum_sb;
211-
neighbors_indices[u + 1] = i0;
211+
(*neighbors_indices)[u] = inds_dists[j].first + sum_sb;
212+
(*neighbors_indices)[u + 1] = i0;
212213
u += 2;
213214
}
214215
}

‎csrc/cpu/utils/neighbors.h‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,12 @@ using namespace std;
1111

1212
template<typename scalar_t>
1313
int nanoflann_neighbors(vector<scalar_t>& queries, vector<scalar_t>& supports,
14-
vector<long>& neighbors_indices, float radius, int dim, int max_num);
14+
vector<long>& neighbors_indices, double radius, int dim, int64_t max_num);
1515

1616
template<typename scalar_t>
1717
int batch_nanoflann_neighbors (vector<scalar_t>& queries,
1818
vector<scalar_t>& supports,
1919
vector<long>& q_batches,
2020
vector<long>& s_batches,
2121
vector<long>& neighbors_indices,
22-
float radius, int dim, int max_num);
22+
double radius, int dim, int64_t max_num);

‎test/test_radius.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import pytest
44
import torch
55
from torch_cluster import radius, radius_graph
6-
76
from .utils import grad_dtypes, devices, tensor
87

98

@@ -115,8 +114,8 @@ def test_radius_graph_pointnet_small(dtype, device):
115114

116115
row, col = radius_graph(x, r=0.2, flow='source_to_target', batch=batch)
117116

118-
edges = set([(i, j) for (i, j) in zip(row.cpu().numpy(),
119-
col.cpu().numpy())])
117+
edges = set([(i, j) for (i, j) in zip(list(row.cpu().numpy()),
118+
list(col.cpu().numpy()))])
120119

121120
truth_row = [10, 11, 7, 9, 9, 1, 9, 1, 6, 7, 0, 11, 0, 10, 15, 12, 20, 16,
122121
34, 31, 44, 43, 42, 41]
@@ -404,8 +403,8 @@ def test_radius_graph_pointnet_medium(dtype, device):
404403

405404
row, col = radius_graph(x, r=0.2, flow='source_to_target', batch=batch)
406405

407-
edges = set([(i, j) for (i, j) in zip(row.cpu().numpy(),
408-
col.cpu().numpy())])
406+
edges = set([(i, j) for (i, j) in zip(list(row.cpu().numpy()),
407+
list(col.cpu().numpy()))])
409408

410409
truth_row = [6, 27, 17, 31, 3, 23, 62, 2, 14, 23, 36, 38, 62, 15, 0, 11,
411410
27, 29, 50, 49, 54, 56, 12, 61, 16, 21, 24, 39, 6, 27, 29,
@@ -573,8 +572,8 @@ def test_radius_graph_ndim(dtype, device):
573572

574573
row, col = radius_graph(x, r=4.4, flow='source_to_target', batch=batch)
575574

576-
edges = set([(i, j) for (i, j) in zip(row.cpu().numpy(),
577-
col.cpu().numpy())])
575+
edges = set([(i, j) for (i, j) in zip(list(row.cpu().numpy()),
576+
list(col.cpu().numpy()))])
578577

579578
truth_row = [2, 3, 2, 3, 0, 1, 3, 4, 0, 1, 2, 4, 2, 3, 6, 7, 9, 10, 5, 7,
580579
8, 9, 10, 5, 6, 10, 6, 5, 6, 10, 5, 6, 7, 9, 13, 11, 16, 17,

0 commit comments

Comments
 (0)