Nearest neighbour correspondence

[nuoboy]

I want to find the nearest neighbour correspondence between two tensors for a cost function implementation. Say I have two pointclouds and I want to find the nearest point in other point cloud from all points in the current pointcloud. Is there anyway to do this efficiently in pytorch operations alone?

[nosmokingsurfer]

@nuoboy Hello,

I think, the most naive way is to count all pairwise distances and apply torch.argmin along desired dim:

import torch

# two clouds of different size
N = 100
M = 150

#generating point clouds
cloud_1 = torch.randn(N,3)
cloud_2 = torch.randn(M,3)

# calculating all distance betwean all possible pairs of points
# shape of dist matrix will be (N,M)
dist = (torch.linalg.norm(cloud_1.transpose(0,1)[None,...] - cloud_2[...,None],axis=1)).transpose(0,1)

# by picking dim value selec to which cloud we look for matches
# if we want to find closest for evey point from cloud_1, then put dim = 1
cloud_1_closest = torch.argmin(dist, dim = 1)

# matches will have a pairs (i,j) where i - index of point in cloud_1, j - index of point in cloud_2
matches = [(k,v.item()) for k,v in enumerate(cloud_1_closest)]

print(matches)

[luisenp]

I'm not aware of any out of the box solution for doing this in pytorch. This thread seems relevant.