У меня есть набор данных облака точек. Я хочу вычислить точки KNN для каждой точки в форме и выбрать их. Я думаю, что я могу сначала вычислить матрицу расстояний и использовать pytorch api torch.topk, чтобы получить k верхний по величине индекс и значения. Однако, согласно документации к pytorch, мы можем предоставить только 1-D тензор, содержащий индексы для индексации к api torch.index_select. Это означает, что я должен использовать цикл для выбора этих точек. Есть ли более эффективный способ сделать это?
Ниже приведены две функции, которые я использую для вычисления PCA для каждой точки в одной форме.
def get_distance_matrix(p1,p2):
p1 = p1.unsqueeze(1)
p2 = p2.unsqueeze(1)
p1 = p1.repeat(1, p2.size(2), 1, 1)
p1 = p1.transpose(1, 2)
p2 = p2.repeat(1, p1.size(1), 1, 1)
dist = torch.add(p1, torch.neg(p2))
dist = torch.norm(dist, 2, dim = 3)
return dist
NUM_LOCAL_POINTS = 16
# estimate each (index th) point normal by using pca
def compute_pca(points, distance_matrix, point_index, batch_index):
distance_vector = distance_matrix[batch_index, point_index] # get the index row from distance matrix
point_vector = points[batch_index, point_index] # get the point from point clouds
values, indices = torch.topk(1.0/distance_vector, NUM_LOCAL_POINTS, dim=0) # select knn points
# create the correct indices shape to slice the points
local_patch = torch.index_select(points[batch_index], 0, indices) - point_vector
# compute the pca and the return eigen_value is sotred in non-descreaing order
local_patch_T = torch.transpose(local_patch, 1, 0)
eigen_values, eigen_vectors = torch.eig(torch.matmul(local_patch_T, local_patch), True)
# find the smallest eigen_value's eigen_vector
# print("Size of point vector is ", point_vector.size())
min_index = torch.argmin(eigen_values[:, 0], 0)
return eigen_vectors[:, min_index]# return the smallest eigen_vectors