'ball_query_2 is working'

Author: tchaton

cuda/src/ball_query.cpp

@@ -69,6 +69,7 @@ std::pair<at::Tensor, at::Tensor> ball_query_partial_dense(at::Tensor x,
 
 	at::Tensor idx = torch::full({x.size(0), nsample}, y.size(0),
 				      at::device(x.device()).dtype(at::ScalarType::Long));
+	
 	at::Tensor dist = torch::full({x.size(0), nsample}, -1,
 				      at::device(x.device()).dtype(at::ScalarType::Float));
 
@@ -78,15 +79,11 @@ std::pair<at::Tensor, at::Tensor> ball_query_partial_dense(at::Tensor x,
 				cudaMemcpyDeviceToHost);
 	auto batch_size = batch_sizes[0] + 1;
 
-	std::cout << batch_x << std::endl;
 	batch_x = degree(batch_x, batch_size);
 	batch_x = at::cat({at::zeros(1, batch_x.options()), batch_x.cumsum(0)}, 0);
 	batch_y = degree(batch_y, batch_size);
 	batch_y = at::cat({at::zeros(1, batch_y.options()), batch_y.cumsum(0)}, 0);
 
-	std::cout << batch_x << std::endl;
-	std::cout << batch_y << std::endl;
-
 	if (x.type().is_cuda()) {
 		query_ball_point_kernel_partial_wrapper(batch_size,
 							x.size(0),

cuda/src/ball_query_gpu.cu

@@ -69,16 +69,13 @@ __global__ void query_ball_point_kernel_partial_dense(int size_x,
 	float radius2 = radius * radius;
 
 	for (ptrdiff_t n_x = start_idx_x + idx; n_x < end_idx_x; n_x += THREADS) {
-		printf("n_x: %d \n", n_x);
-
 		int64_t count = 0;
 		for (ptrdiff_t n_y = start_idx_y; n_y < end_idx_y; n_y++) {
 			float dist = 0;
 			for (ptrdiff_t d = 0; d < 3; d++) {
 				dist += (x[n_x * 3 + d] - y[n_y * 3 + d]) *
 					(x[n_x * 3 + d] - y[n_y * 3 + d]);
 			}
-			printf("n_x: %d, n_y: %d \n", n_x, n_y);
 			if(dist <= radius2){
 				idx_out[n_x * nsample + count] = n_y;
 				dist_out[n_x * nsample + count] = dist;

test/test_ballquerry_partial.py

@@ -13,8 +13,6 @@ def test_simple_gpu(self):
         batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long().cuda()
         batch_y = torch.from_numpy(np.asarray([0])).long().cuda()
 
-        print(x.shape, y.shape, batch_x.shape, batch_y.shape)
-
         batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long().cuda()
         batch_y = torch.from_numpy(np.asarray([0])).long().cuda()
 
@@ -23,11 +21,11 @@ def test_simple_gpu(self):
         idx = idx.detach().cpu().numpy()
         dist2 = dist2.detach().cpu().numpy()
 
-        print(idx)
-        print(dist2)
+        idx_answer = np.asarray([[1, 1], [0, 1], [1, 1], [1, 1]])
+        dist2_answer = np.asarray([[-1, -1], [0.01, -1], [-1, -1], [-1, -1]]).astype(np.float32)
 
-        #npt.assert_array_equal(idx, )
-        #npt.assert_array_equal(dist2, np.asarray([[0.3, 0.1]]))
+        npt.assert_array_almost_equal(idx, idx_answer)
+        npt.assert_array_almost_equal(dist2, dist2_answer)