def test_speed(self):
start = time.time()
a = torch.randn(50000, 3).to(torch.float)
b = torch.randn(10000, 3).to(torch.float)
batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] + [1 for i in range(a.shape[0] // 2, a.shape[0])])
batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] + [1 for i in range(b.shape[0] // 2, b.shape[0])])
R = 1
samples = 50
idx, dist = ball_query(
R,
samples,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=True,
)
idx1, dist = ball_query(
R,
samples,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=True,
)
print(time.time() - start)
torch.testing.assert_allclose(idx1, idx)
self.assertEqual(idx.shape[0], b.shape[0])
self.assertEqual(dist.shape[0], b.shape[0])
self.assertLessEqual(idx.max().item(), len(batch_a))
def test_random_cpu(self):
a = torch.randn(100, 3).to(torch.float)
b = torch.randn(50, 3).to(torch.float)
batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] +
[1 for i in range(a.shape[0] // 2, a.shape[0])])
batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] +
[1 for i in range(b.shape[0] // 2, b.shape[0])])
R = 1
idx, dist = ball_query(R,
15,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=True)
idx1, dist = ball_query(R,
15,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=True)
torch.testing.assert_allclose(idx1, idx)
with self.assertRaises(AssertionError):
idx, dist = ball_query(R,
15,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=False)
idx1, dist = ball_query(R,
15,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=False)
torch.testing.assert_allclose(idx1, idx)
self.assertEqual(idx.shape[0], b.shape[0])
self.assertEqual(dist.shape[0], b.shape[0])
self.assertLessEqual(idx.max().item(), len(batch_a))
# Comparison to see if we have the same result
tree = KDTree(a.detach().numpy())
idx3_sk = tree.query_radius(b.detach().numpy(), r=R)
i = np.random.randint(len(batch_b))
for p in idx[i].detach().numpy():
if p >= 0 and p < len(batch_a):
assert p in idx3_sk[i]
def test_simple_cpu(self):
a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]],
[[0, 0, 0], [1, 0, 0], [2, 0, 0]]]).to(torch.float)
b = torch.tensor([[[0, 0, 0]], [[3, 0, 0]]]).to(torch.float)
idx, dist = ball_query(1.01, 2, a, b, sort=True)
torch.testing.assert_allclose(idx, torch.tensor([[[0, 1]], [[2, 2]]]))
torch.testing.assert_allclose(
dist,
torch.tensor([[[0, 1]], [[1, -1]]]).float())
a = torch.tensor([[[0, 0, 0], [1, 0, 0], [1, 1, 0]]]).to(torch.float)
idx, dist = ball_query(1.01, 3, a, a, sort=True)
torch.testing.assert_allclose(
idx, torch.tensor([[[0, 1, 0], [1, 0, 2], [2, 1, 2]]]))
def test_simple_gpu(self):
x = torch.tensor([[10, 0, 0], [0.1, 0, 0], [10, 0, 0],
[0.1, 0, 0]]).to(torch.float).cuda()
y = torch.tensor([[0, 0, 0]]).to(torch.float).cuda()
batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long().cuda()
batch_y = torch.from_numpy(np.asarray([0])).long().cuda()
batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long().cuda()
batch_y = torch.from_numpy(np.asarray([0])).long().cuda()
idx, dist2 = ball_query(1.0,
2,
x,
y,
mode="PARTIAL_DENSE",
batch_x=batch_x,
batch_y=batch_y)
idx = idx.detach().cpu().numpy()
dist2 = dist2.detach().cpu().numpy()
idx_answer = np.asarray([[1, -1]])
dist2_answer = np.asarray([[0.0100, -1.0000]]).astype(np.float32)
npt.assert_array_almost_equal(idx, idx_answer)
npt.assert_array_almost_equal(dist2, dist2_answer)
def test_random_gpu(self):
a = torch.randn(100, 3).to(torch.float).cuda()
b = torch.randn(50, 3).to(torch.float).cuda()
batch_a = torch.tensor([0 for i in range(a.shape[0] // 2)] +
[1 for i in range(a.shape[0] //
2, a.shape[0])]).cuda()
batch_b = torch.tensor([0 for i in range(b.shape[0] // 2)] +
[1 for i in range(b.shape[0] //
2, b.shape[0])]).cuda()
R = 1
idx, dist = ball_query(
R,
15,
a,
b,
mode="PARTIAL_DENSE",
batch_x=batch_a,
batch_y=batch_b,
sort=False,
)
# Comparison to see if we have the same result
tree = KDTree(a.cpu().detach().numpy())
idx3_sk = tree.query_radius(b.cpu().detach().numpy(), r=R)
i = np.random.randint(len(batch_b))
for p in idx[i].cpu().detach().numpy():
if p >= 0 and p < len(batch_a):
assert p in idx3_sk[i]
def test_cpu_gpu_equality(self):
a = torch.randn(5, 1000, 3)
b = torch.randn(5, 500, 3)
res_cpu = ball_query(1, 500, a, b)[0].detach().numpy()
res_cuda = ball_query(1, 500, a.cuda(),
b.cuda())[0].cpu().detach().numpy()
for i in range(b.shape[0]):
for j in range(b.shape[1]):
# Because it is not necessary the same order
assert set(res_cpu[i][j]) == set(res_cuda[i][j])
res_cpu = ball_query(0.01, 500, a, b)[0].detach().numpy()
res_cuda = ball_query(0.01, 500, a.cuda(),
b.cuda())[0].cpu().detach().numpy()
for i in range(b.shape[0]):
for j in range(b.shape[1]):
# Because it is not necessary the same order
assert set(res_cpu[i][j]) == set(res_cuda[i][j])
def _variance_estimator_dense(r, pos, f):
nei_idx = tp.ball_query(r, _MAX_NEIGHBOURS, pos, pos,
sort=True)[0].reshape(pos.shape[0],
-1).long() # [B,N * nei]
f_neighboors = f.gather(1, nei_idx).reshape(f.shape[0], f.shape[1],
-1) # [B,N , nei]
gradient = (f.unsqueeze(-1).repeat(1, 1, f_neighboors.shape[-1]) -
f_neighboors)**2 # [B,N,nei]
return gradient.sum(-1)
def test_simple_gpu(self):
a = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]],
[[0, 0, 0], [1, 0, 0], [2, 0,
0]]]).to(torch.float).cuda()
b = torch.tensor([[[0, 0, 0]], [[3, 0, 0]]]).to(torch.float).cuda()
idx, dist = ball_query(1.01, 2, a, b)
torch.testing.assert_allclose(idx.cpu(),
torch.tensor([[[0, 1]], [[2, 2]]]))
torch.testing.assert_allclose(
dist.cpu(),
torch.tensor([[[0, 1]], [[1, -1]]]).float())
def test_breaks(self):
x = torch.tensor([[10, 0, 0], [0.1, 0, 0], [10, 0, 0],
[10.1, 0, 0]]).to(torch.float)
y = torch.tensor([[0, 0, 0]]).to(torch.float)
batch_x = torch.from_numpy(np.asarray([0, 0, 1, 1])).long()
batch_y = torch.from_numpy(np.asarray([0])).long()
with self.assertRaises(RuntimeError):
idx, dist2 = ball_query(1.0,
2,
x,
y,
mode="PARTIAL_DENSE",
batch_x=batch_x,
batch_y=batch_y)
def find_neighbours(self, x, y, scale_idx=0):
if scale_idx >= self.num_scales:
raise ValueError("Scale %i is out of bounds %i" %
(scale_idx, self.num_scales))
num_neighbours = self._max_num_neighbors[scale_idx]
neighbours = tp.ball_query(self._radius[scale_idx], num_neighbours, x,
y)[0]
if DEBUGGING_VARS["FIND_NEIGHBOUR_DIST"]:
for i in range(neighbours.shape[0]):
start = neighbours[i, :, 0]
valid_neighbours = (neighbours[i, :, 1:] != start.view(
(-1, 1)).repeat(1, num_neighbours - 1)).sum(1) + 1
self._dist_meters[scale_idx].add_valid_neighbours(
valid_neighbours)
return neighbours
def find_neighbours(self, x, y, batch_x=None, batch_y=None):
if self._conv_type == ConvolutionFormat.MESSAGE_PASSING.value:
return radius(x,
y,
self._radius,
batch_x,
batch_y,
max_num_neighbors=self._max_num_neighbors)
elif self._conv_type == ConvolutionFormat.DENSE.value or ConvolutionFormat.PARTIAL_DENSE.value:
return tp.ball_query(self._radius,
self._max_num_neighbors,
x,
y,
mode=self._conv_type,
batch_x=batch_x,
batch_y=batch_y)[0]
else:
raise NotImplementedError
def test_larger_gpu(self):
a = torch.randn(32, 4096, 3).to(torch.float).cuda()
idx, dist = ball_query(1, 64, a, a)
self.assertGreaterEqual(idx.min(), 0)
