def test_invert_neighbors_list_shape_checking(ml): num_points = 3 inp_neighbors_index = np.array([0, 1, 2, 2, 1, 2], dtype=np.int32) inp_neighbors_row_splits = np.array([0, 3, 4, 6], dtype=np.int64) inp_neighbors_attributes = np.array([10, 20, 30, 40, 50, 60], dtype=np.float32) # test the shape checking by passing arrays with wrong rank and/or size with pytest.raises(Exception) as einfo: _ = mltest.run_op(ml, ml.cpu_device, False, ml.ops.invert_neighbors_list, num_points=num_points, inp_neighbors_index=inp_neighbors_index[1:], inp_neighbors_row_splits=inp_neighbors_row_splits, inp_neighbors_attributes=inp_neighbors_attributes) assert 'invalid shape' in str(einfo.value) with pytest.raises(Exception) as einfo: _ = mltest.run_op(ml, ml.cpu_device, False, ml.ops.invert_neighbors_list, num_points=num_points, inp_neighbors_index=inp_neighbors_index[:, np.newaxis], inp_neighbors_row_splits=inp_neighbors_row_splits, inp_neighbors_attributes=inp_neighbors_attributes) assert 'invalid shape' in str(einfo.value) with pytest.raises(Exception) as einfo: _ = mltest.run_op( ml, ml.cpu_device, False, ml.ops.invert_neighbors_list, num_points=num_points, inp_neighbors_index=inp_neighbors_index, inp_neighbors_row_splits=inp_neighbors_row_splits[:, np.newaxis], inp_neighbors_attributes=inp_neighbors_attributes) assert 'invalid shape' in str(einfo.value) with pytest.raises(Exception) as einfo: _ = mltest.run_op( ml, ml.cpu_device, False, ml.ops.invert_neighbors_list, num_points=num_points, inp_neighbors_index=inp_neighbors_index, inp_neighbors_row_splits=inp_neighbors_row_splits, inp_neighbors_attributes=inp_neighbors_attributes[1:]) assert 'invalid shape' in str(einfo.value)
def test_roi_pool(ml): values0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/values0.npy' ) values1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/values1.npy' ) values2 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/values2.npy' ) sampled_pts_num = 512 ans0, ans1 = mltest.run_op(ml, ml.device, True, ml.ops.roi_pool, values0, values1, values2, sampled_pts_num) expected0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/out0.npy' ) expected1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/roi_pool/out1.npy' ) np.testing.assert_equal(ans0, expected0) np.testing.assert_equal(ans1, expected1)
def test_reduce_subarrays_sum_random(seed, dtype, device, ml): rng = np.random.RandomState(seed) values_shape = [rng.randint(100, 200)] values = rng.uniform(0, 10, size=values_shape).astype(dtype) row_splits = [0] for _ in range(rng.randint(1, 10)): row_splits.append( rng.randint(0, values_shape[0] - row_splits[-1]) + row_splits[-1]) row_splits.extend(values_shape) expected_result = [] for start, stop in zip(row_splits, row_splits[1:]): # np.sum correctly handles zero length arrays and returns 0 expected_result.append(np.sum(values[start:stop])) np.array(expected_result, dtype=dtype) row_splits = np.array(row_splits, dtype=np.int64) ans = mltest.run_op(ml, device, True, ml.ops.reduce_subarrays_sum, values=values, row_splits=row_splits) if np.issubdtype(dtype, np.integer): np.testing.assert_equal(ans, expected_result) else: # floating point types np.testing.assert_allclose(ans, expected_result, rtol=1e-5, atol=1e-8)
def test_voxelize_random(ml, point_dtype, ndim, max_voxels, max_points_per_voxel): rng = np.random.RandomState(123) points = rng.rand(rng.randint(0, 10000), ndim).astype(point_dtype) voxel_size = rng.uniform(0.01, 0.1, size=(ndim, )).astype(point_dtype) point_range_min = rng.uniform(0.0, 0.3, size=(ndim, )).astype(point_dtype) point_range_max = rng.uniform(0.7, 1.0, size=(ndim, )).astype(point_dtype) ans = mltest.run_op(ml, ml.device, True, ml.ops.voxelize, points, voxel_size, point_range_min, point_range_max, max_points_per_voxel=max_points_per_voxel, max_voxels=max_voxels) voxels = convert_output_to_voxel_dict(ans) voxels_reference = voxelize_python( points, voxel_size, point_range_min, point_range_max, max_points_per_voxel=max_points_per_voxel, max_voxels=max_voxels) assert_equal_voxel_dicts(voxels, ref=voxels_reference)
def sparse_conv_transpose_filter(filters): return mltest.run_op(ml, ml.device, True, ml.ops.sparse_conv_transpose, filters, inp_importance, y_arr, neighbors_index, neighbors_importance_sum, neighbors_row_splits, inv_neighbors_index, inv_neighbors_kernel_index, inv_neighbors_importance, inv_neighbors_row_splits, **conv_attrs)
def test_fixed_radius_search_empty_point_sets(ml): rng = np.random.RandomState(123) dtype = np.float32 radius = 1 hash_table_size_factor = 1 / 64 # no query points points = rng.random(size=(100, 3)).astype(dtype) queries = rng.random(size=(0, 3)).astype(dtype) layer = ml.layers.FixedRadiusSearch(return_distances=True) ans = mltest.run_op( ml, ml.device, True, layer, points, queries=queries, radius=radius, hash_table_size_factor=hash_table_size_factor, ) assert ans.neighbors_index.shape == (0, ) assert ans.neighbors_row_splits.shape == (1, ) assert ans.neighbors_distance.shape == (0, ) # no input points points = rng.random(size=(0, 3)).astype(dtype) queries = rng.random(size=(100, 3)).astype(dtype) ans = mltest.run_op( ml, ml.device, True, layer, points, queries=queries, radius=radius, hash_table_size_factor=hash_table_size_factor, ) assert ans.neighbors_index.shape == (0, ) assert ans.neighbors_row_splits.shape == (101, ) np.testing.assert_array_equal(np.zeros_like(ans.neighbors_row_splits), ans.neighbors_row_splits) assert ans.neighbors_distance.shape == (0, )
def sparse_conv_transpose_infeats(inp_features): return mltest.run_op(ml, ml.device, True, ml.ops.sparse_conv_transpose, filters.transpose([0, 2, 1]), inp_importance, inp_features, neighbors_index, neighbors_importance_sum, neighbors_row_splits, inv_neighbors_index, inv_neighbors_kernel_index, inv_neighbors_importance, inv_neighbors_row_splits, **conv_attrs)
def test_knn_search_empty_point_sets(ml): rng = np.random.RandomState(123) device = mltest.cpu_device dtype = np.float32 # no query points points = rng.random(size=(100, 3)).astype(dtype) queries = rng.random(size=(0, 3)).astype(dtype) k = rng.randint(1, 11) layer = ml.layers.KNNSearch(return_distances=True) ans = mltest.run_op( ml, device, layer, True, points, queries=queries, k=k, ) assert ans.neighbors_index.shape == (0, ) assert ans.neighbors_row_splits.shape == (1, ) assert ans.neighbors_distance.shape == (0, ) # no input points points = rng.random(size=(0, 3)).astype(dtype) queries = rng.random(size=(100, 3)).astype(dtype) layer = ml.layers.KNNSearch(return_distances=True) ans = mltest.run_op( ml, device, layer, True, points, queries=queries, k=k, ) assert ans.neighbors_index.shape == (0, ) assert ans.neighbors_row_splits.shape == (101, ) np.testing.assert_array_equal(np.zeros_like(ans.neighbors_row_splits), ans.neighbors_row_splits) assert ans.neighbors_distance.shape == (0, )
def test_radius_search_empty_point_sets(ml): rng = np.random.RandomState(123) dtype = np.float32 # no query points points = rng.random(size=(100, 3)).astype(dtype) queries = rng.random(size=(0, 3)).astype(dtype) radii = rng.uniform(0.1, 0.3, size=(0, )).astype(dtype) layer = ml.layers.RadiusSearch(return_distances=True) ans = mltest.run_op( ml, ml.device, True, layer, points, queries=queries, radii=radii, ) assert ans.neighbors_index.shape == (0, ) assert ans.neighbors_row_splits.shape == (1, ) assert ans.neighbors_distance.shape == (0, ) # no input points points = rng.random(size=(0, 3)).astype(dtype) queries = rng.random(size=(100, 3)).astype(dtype) radii = rng.uniform(0.1, 0.3, size=(100, )).astype(dtype) ans = mltest.run_op( ml, ml.device, True, layer, points, queries=queries, radii=radii, ) assert ans.neighbors_index.shape == (0, ) assert ans.neighbors_row_splits.shape == (101, ) np.testing.assert_array_equal(np.zeros_like(ans.neighbors_row_splits), ans.neighbors_row_splits) assert ans.neighbors_distance.shape == (0, )
def test_fixed_radius_search(dtype, device, ml, num_points_queries, radius, hash_table_size_factor, metric, ignore_query_point, return_distances): # skip dtype not supported on GPU if 'GPU' in device and not dtype in gpu_dtypes: return rng = np.random.RandomState(123) num_points, num_queries = num_points_queries points = rng.random(size=(num_points, 3)).astype(dtype) if ignore_query_point: queries = points else: queries = rng.random(size=(num_queries, 3)).astype(dtype) # kd tree for computing the ground truth tree = cKDTree(points, copy_data=True) p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric] gt_neighbors_index = tree.query_ball_point(queries, radius, p=p_norm) layer = ml.layers.FixedRadiusSearch(metric=metric, ignore_query_point=ignore_query_point, return_distances=return_distances) ans = mltest.run_op( ml, device, True, layer, points, queries=queries, radius=radius, hash_table_size_factor=hash_table_size_factor, ) for i, q in enumerate(queries): # check neighbors start = ans.neighbors_row_splits[i] end = ans.neighbors_row_splits[i + 1] q_neighbors_index = ans.neighbors_index[start:end] gt_set = set(gt_neighbors_index[i]) if ignore_query_point: gt_set.remove(i) assert gt_set == set(q_neighbors_index) # check distances if return_distances: q_neighbors_dist = ans.neighbors_distance[start:end] for j, dist in zip(q_neighbors_index, q_neighbors_dist): if metric == 'L2': gt_dist = np.sum((q - points[j])**2) else: gt_dist = np.linalg.norm(q - points[j], ord=p_norm) np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)
def test_voxelize_simple(ml, point_dtype, point_range_max, max_voxels, max_points_per_voxel): # global iteration # if iteration > 0: # return # iteration += 1 # if 'GPU' in ml.device: # return # yapf: disable points = np.array([ # 2 points in voxel [0.5, 0.5, 0.5], [0.7, 0.2, 0.3], # new batch # 1 point in two voxels [0.7, 0.5, 0.9], [10.7, 10.2, 10.3], # 2 points in another voxel [1.4, 1.5, 1.4], [1.7, 1.2, 1.3], ], dtype=point_dtype) row_splits = np.array([0, 2, 4, 6], dtype=np.int64) # row_splits = np.array([0,6], dtype=np.int64) # yapf: enable voxel_size = np.array([1.0, 1.1, 1.2], dtype=point_dtype) point_range_min = np.zeros((3, ), dtype=point_dtype) point_range_max = np.array(point_range_max, dtype=point_dtype) ans = mltest.run_op(ml, ml.device, True, ml.ops.voxelize, points, row_splits, voxel_size, point_range_min, point_range_max, max_points_per_voxel=max_points_per_voxel, max_voxels=max_voxels) voxels = convert_output_to_voxel_dict(ans) voxels_reference = voxelize_python( points, row_splits, voxel_size, point_range_min, point_range_max, max_points_per_voxel=max_points_per_voxel, max_voxels=max_voxels) # print(voxels) # print(voxels_reference) # print('max_voxels', max_voxels, 'max_points_per_voxel', # max_points_per_voxel) assert_equal_voxel_dicts(voxels, ref_b=voxels_reference)
def test_voxel_pooling_empty_point_set(ml, pos_dtype, feat_dtype, position_fn, feature_fn): points = np.zeros(shape=[0, 3], dtype=pos_dtype) features = np.zeros(shape=[0, 5], dtype=feat_dtype) voxel_size = 1 ans = mltest.run_op(ml, ml.device, True, ml.ops.voxel_pooling, points, features, voxel_size, position_fn, feature_fn) np.testing.assert_array_equal(points, ans.pooled_positions) np.testing.assert_array_equal(features, ans.pooled_features)
def test_furthest_point_sampling(ml): values = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/sampling/values.npy' ) samples = 4096 ans = mltest.run_op(ml, ml.device, True, ml.ops.furthest_point_sampling, values, samples) expected = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/sampling/out.npy' ) np.testing.assert_equal(ans, expected)
def test_nms_empty(ml): boxes = np.zeros((0, 5), dtype=np.float32) scores = np.array([], dtype=np.float32) nms_overlap_thresh = 0.7 keep_indices_ref = np.array([]).astype(np.int64) keep_indices = mltest.run_op(ml, ml.device, True, ml.ops.nms, boxes, scores, nms_overlap_thresh=nms_overlap_thresh) np.testing.assert_equal(keep_indices, keep_indices_ref) assert keep_indices.dtype == keep_indices_ref.dtype
def test_group_pts(ml): values0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/values0.npy' ) values1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/values1.npy' ) ans = mltest.run_op(ml, ml.device, True, ml.ops.group_points, values0, values1) expected = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/group_pts/out.npy' ) np.testing.assert_equal(ans, expected)
def test_cublas_matmul(ml): # This test checks if calling cublas functionality from open3d and the ml framework works. rng = np.random.RandomState(123) n = 20 arr = rng.rand(n, n).astype(np.float32) # do matmul with open3d A = o3d.core.Tensor.from_numpy(arr).cuda() B = A @ A # now use the ml framework cublas C = mltest.run_op(ml, ml.device, True, ml.module.matmul, arr, arr) np.testing.assert_allclose(B.cpu().numpy(), C)
def conv_filter(filters): return mltest.run_op(ml, ml.device, True, ml.ops.continuous_conv, filters=filters, out_positions=out_positions, extents=extent, offset=offset, inp_positions=inp_positions, inp_features=inp_features, inp_importance=inp_importance, neighbors_index=neighbors_index, neighbors_importance=neighbors_importance, neighbors_row_splits=neighbors_row_splits, **conv_attrs)
def test_ragged_to_dense(dtype, ml): values = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=dtype) row_splits = np.array([0, 2, 4, 4, 5, 12, 13], dtype=np.int64) out_col_size = 4 default_value = np.array(-1, dtype=dtype) ans = mltest.run_op(ml, ml.device, True, ml.ops.ragged_to_dense, values, row_splits, out_col_size, default_value) expected = np.full((row_splits.shape[0] - 1, out_col_size), default_value) for i in range(row_splits.shape[0] - 1): for j, value_idx in enumerate(range(row_splits[i], row_splits[i + 1])): if j < expected.shape[1]: expected[i, j] = values[value_idx] np.testing.assert_equal(ans, expected)
def test_query_pts(ml): values0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/query_pts/values0.npy' ) values1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/query_pts/values1.npy' ) sample = 16 radius = 0.1 ans = mltest.run_op(ml, ml.device, True, ml.ops.ball_query, values0, values1, radius, sample) expected = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/query_pts/out.npy' ) np.testing.assert_equal(ans, expected)
def test_three_interp(ml): values0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/values0.npy' ) values1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/values1.npy' ) values2 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/values2.npy' ) ans = mltest.run_op(ml, ml.device, True, ml.ops.three_interpolate, values0, values1, values2) expected = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_interp/out.npy' ) np.testing.assert_equal(ans, expected)
def test_three_nn(ml): values0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/values0.npy' ) values1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/values1.npy' ) ans0, ans1 = mltest.run_op(ml, ml.device, True, ml.ops.three_nn, values0, values1) expected0 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/out0.npy' ) expected1 = mltest.fetch_numpy( 'https://storage.googleapis.com/isl-datasets/open3d-dev/test/ml_ops/data/three_nn/out1.npy' ) np.testing.assert_equal(ans0, expected0) np.testing.assert_equal(ans1, expected1)
def conv_transpose_infeats(inp_features): return mltest.run_op( ml, ml.device, True, ml.ops.continuous_conv_transpose, filters=filters.transpose([0, 1, 2, 4, 3]), out_positions=inp_positions, out_importance=inp_importance, extents=extent, offset=offset, inp_positions=out_positions, inp_features=inp_features, inp_neighbors_index=neighbors_index, inp_neighbors_importance_sum=neighbors_importance_sum, inp_neighbors_row_splits=neighbors_row_splits, neighbors_index=inverted_neighbors_index, neighbors_importance=inverted_neighbors_importance, neighbors_row_splits=inverted_neighbors_row_splits, **conv_attrs)
def test_ragged_to_dense_more_dims(dtype, ml): values = np.array([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7], [8, 8], [9, 9], [10, 10], [11, 11], [12, 12]], dtype=dtype) row_splits = np.array([0, 2, 4, 4, 5, 12, 13], dtype=np.int64) out_col_size = 4 default_value = np.array([-1, -1], dtype=dtype) ans = mltest.run_op(ml, ml.device, True, ml.ops.ragged_to_dense, values, row_splits, out_col_size, default_value) expected = np.full(( row_splits.shape[0] - 1, out_col_size, ) + default_value.shape, default_value) for i in range(row_splits.shape[0] - 1): for j, value_idx in enumerate(range(row_splits[i], row_splits[i + 1])): if j < expected.shape[1]: expected[i, j] = values[value_idx] np.testing.assert_equal(ans, expected)
def test_nms(ml): boxes = np.array([[15.0811, -7.9803, 15.6721, -6.8714, 0.5152], [15.1166, -7.9261, 15.7060, -6.8137, 0.6501], [15.1304, -7.8129, 15.7069, -6.8903, 0.7296], [15.2050, -7.8447, 15.8311, -6.7437, 1.0506], [15.1343, -7.8136, 15.7121, -6.8479, 1.0352], [15.0931, -7.9552, 15.6675, -7.0056, 0.5979]], dtype=np.float32) scores = np.array([3, 1.1, 5, 2, 1, 0], dtype=np.float32) nms_overlap_thresh = 0.7 keep_indices_ref = np.array([2, 3, 5]).astype(np.int64) keep_indices = mltest.run_op(ml, ml.device, True, ml.ops.nms, boxes, scores, nms_overlap_thresh=nms_overlap_thresh) np.testing.assert_equal(keep_indices, keep_indices_ref) assert keep_indices.dtype == keep_indices_ref.dtype
def test_reduce_subarrays_sum_zero_length_values(ml): rng = np.random.RandomState(1) shape = [rng.randint(100, 200)] values = np.array([], dtype=np.float32) row_splits = [0] for _ in range(rng.randint(1, 10)): row_splits.append( rng.randint(0, shape[0] - row_splits[-1]) + row_splits[-1]) row_splits.extend(shape) row_splits = np.array(row_splits, dtype=np.int64) ans = mltest.run_op(ml, ml.device, True, ml.ops.reduce_subarrays_sum, values=values, row_splits=row_splits) assert ans.shape == values.shape assert ans.dtype == values.dtype
def test_voxelize_random(ml, point_dtype, ndim, batch_size, max_voxels, max_points_per_voxel): rng = np.random.RandomState(123) points = rng.rand(rng.randint(0, 10000), ndim).astype(point_dtype) row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64) for i in range(batch_size): row_splits[i + 1] = rng.randint( points.shape[0] // batch_size) + row_splits[i] points = points[:row_splits[batch_size]] voxel_size = rng.uniform(0.01, 0.1, size=(ndim, )).astype(point_dtype) point_range_min = rng.uniform(0.0, 0.3, size=(ndim, )).astype(point_dtype) point_range_max = rng.uniform(0.7, 1.0, size=(ndim, )).astype(point_dtype) ans = mltest.run_op(ml, ml.device, True, ml.ops.voxelize, points, row_splits, voxel_size, point_range_min, point_range_max, max_points_per_voxel=max_points_per_voxel, max_voxels=max_voxels) voxels = convert_output_to_voxel_dict(ans) voxels_reference = voxelize_python( points, row_splits, voxel_size, point_range_min, point_range_max, max_points_per_voxel=max_points_per_voxel, max_voxels=max_voxels) assert_equal_voxel_dicts(voxels, ref_b=voxels_reference)
def test_ragged_to_dense_random(dtype, ml, seed): rng = np.random.RandomState(seed) values = rng.random(size=(10000, )).astype(dtype) row_splits = [0] while row_splits[-1] < values.shape[0]: row_splits.append(row_splits[-1] + rng.randint(0, 10)) row_splits[-1] = values.shape[0] row_splits = np.array(row_splits, dtype=np.int64) out_col_size = rng.randint(1, 37) default_value = np.array(-1, dtype=dtype) ans = mltest.run_op(ml, ml.device, True, ml.ops.ragged_to_dense, values, row_splits, out_col_size, default_value) expected = np.full((row_splits.shape[0] - 1, out_col_size), default_value) for i in range(row_splits.shape[0] - 1): for j, value_idx in enumerate(range(row_splits[i], row_splits[i + 1])): if j < expected.shape[1]: expected[i, j] = values[value_idx] np.testing.assert_equal(ans, expected)
def test_invert_neighbors_list(dtype, attributes, ml): # yapf: disable # define connectivity for 3 query points and 3 input points num_points = 3 edges = np.array( [ [0, 0], [0, 1], [0, 2], # 3 neighbors [1, 2], # 1 neighbors [2, 1], [2, 2], # 2 neighbors ], dtype=np.int32) # the neighbors_index is the second column neighbors_index = edges[:, 1] # exclusive prefix sum of the number of neighbors neighbors_row_splits = np.array([0, 3, 4, edges.shape[0]], dtype=np.int64) if attributes == 'scalar': neighbors_attributes = np.array([ 10, 20, 30, 40, 50, 60, ], dtype=dtype) elif attributes == 'none': neighbors_attributes = np.array([], dtype=dtype) elif attributes == 'multidim': neighbors_attributes = np.array([ [10, 1], [20, 2], [30, 3], [40, 4], [50, 5], [60, 6], ], dtype=dtype) # yapf: enable ans = mltest.run_op(ml, ml.device, True, ml.ops.invert_neighbors_list, num_points=num_points, inp_neighbors_index=neighbors_index, inp_neighbors_row_splits=neighbors_row_splits, inp_neighbors_attributes=neighbors_attributes) expected_neighbors_row_splits = [0, 1, 3, edges.shape[0]] np.testing.assert_equal(ans.neighbors_row_splits, expected_neighbors_row_splits) # checking the neighbors_index is more complicated because the order # of the neighbors for each query point is not defined. expected_neighbors_index = [ set([0]), set([0, 2]), set([0, 1, 2]), ] for i, expected_neighbors_i in enumerate(expected_neighbors_index): start = ans.neighbors_row_splits[i] end = ans.neighbors_row_splits[i + 1] neighbors_i = set(ans.neighbors_index[start:end]) assert neighbors_i == expected_neighbors_i if neighbors_attributes.shape == (0, ): # if the input is a zero length vector then the returned attributes # vector also must be a zero length vector assert ans.neighbors_attributes.shape == (0, ) else: # check if the attributes are still associated with the same edge edge_attr_map = { tuple(k): v for k, v in zip(edges, neighbors_attributes) } for i, _ in enumerate(expected_neighbors_index): start = ans.neighbors_row_splits[i] end = ans.neighbors_row_splits[i + 1] # neighbors and attributes for point i neighbors_i = ans.neighbors_index[start:end] attributes_i = ans.neighbors_attributes[start:end] for j, attr in zip(neighbors_i, attributes_i): key = (j, i) np.testing.assert_equal(attr, edge_attr_map[key])
def test_knn_search(dtype, ml, num_points_queries, metric, ignore_query_point, return_distances): rng = np.random.RandomState(123) device = mltest.cpu_device num_points, num_queries = num_points_queries points = rng.random(size=(num_points, 3)).astype(dtype) if ignore_query_point: queries = points else: queries = rng.random(size=(num_queries, 3)).astype(dtype) k = rng.randint(1, 11) # kd tree for computing the ground truth tree = cKDTree(points, copy_data=True) p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric] if k > num_points: gt_neighbors_index = [tree.query(q, k, p=p_norm) for q in queries] gt_neighbors_index = [ idxs[np.isfinite(dists)] for dists, idxs in gt_neighbors_index ] else: gt_neighbors_index = [tree.query(q, k, p=p_norm)[1] for q in queries] layer = ml.layers.KNNSearch(metric=metric, ignore_query_point=ignore_query_point, return_distances=return_distances) ans = mltest.run_op( ml, device, layer, True, points, queries=queries, k=k, ) for i, q in enumerate(queries): # check neighbors start = ans.neighbors_row_splits[i] end = ans.neighbors_row_splits[i + 1] q_neighbors_index = ans.neighbors_index[start:end] if k == 1: gt_set = set([gt_neighbors_index[i]]) else: gt_set = set(gt_neighbors_index[i]) if ignore_query_point: gt_set.remove(i) assert gt_set == set(q_neighbors_index) # check distances if return_distances: q_neighbors_dist = ans.neighbors_distance[start:end] for j, dist in zip(q_neighbors_index, q_neighbors_dist): if metric == 'L2': gt_dist = np.sum((q - points[j])**2) else: gt_dist = np.linalg.norm(q - points[j], ord=p_norm) np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)
def test_knn_search_batches(ml, batch_size): device = mltest.cpu_device dtype = np.float32 metric = 'L2' p_norm = {'L1': 1, 'L2': 2, 'Linf': np.inf}[metric] ignore_query_point = False return_distances = True rng = np.random.RandomState(123) # create array defining start and end of each batch points_row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64) queries_row_splits = np.zeros(shape=(batch_size + 1, ), dtype=np.int64) for i in range(batch_size): points_row_splits[i + 1] = rng.randint(15) + points_row_splits[i] queries_row_splits[i + 1] = rng.randint(15) + queries_row_splits[i] num_points = points_row_splits[-1] num_queries = queries_row_splits[-1] points = rng.random(size=(num_points, 3)).astype(dtype) if ignore_query_point: queries = points queries_row_splits = points_row_splits else: queries = rng.random(size=(num_queries, 3)).astype(dtype) k = rng.randint(1, 11) # kd tree for computing the ground truth gt_neighbors_index = [] for i in range(batch_size): points_i = points[points_row_splits[i]:points_row_splits[i + 1]] queries_i = queries[queries_row_splits[i]:queries_row_splits[i + 1]] tree = cKDTree(points_i, copy_data=True) if k > points_i.shape[0]: tmp = [tree.query(q, k, p=p_norm) for q in queries_i] tmp = [ list(idxs[np.isfinite(dists)] + points_row_splits[i]) for dists, idxs in tmp ] else: tmp = [ list(tree.query(q, k, p=p_norm)[1] + points_row_splits[i]) for q in queries_i ] gt_neighbors_index.extend(tmp) layer = ml.layers.KNNSearch(metric=metric, ignore_query_point=ignore_query_point, return_distances=return_distances) ans = mltest.run_op(ml, device, layer, True, points, queries=queries, k=k, points_row_splits=points_row_splits, queries_row_splits=queries_row_splits) for i, q in enumerate(queries): # check neighbors start = ans.neighbors_row_splits[i] end = ans.neighbors_row_splits[i + 1] q_neighbors_index = ans.neighbors_index[start:end] if k == 1: gt_set = set([gt_neighbors_index[i]]) else: gt_set = set(gt_neighbors_index[i]) if ignore_query_point: gt_set.remove(i) assert gt_set == set(q_neighbors_index) # check distances if return_distances: q_neighbors_dist = ans.neighbors_distance[start:end] for j, dist in zip(q_neighbors_index, q_neighbors_dist): if metric == 'L2': gt_dist = np.sum((q - points[j])**2) else: gt_dist = np.linalg.norm(q - points[j], ord=p_norm) np.testing.assert_allclose(dist, gt_dist, rtol=1e-7, atol=1e-8)