python - 内存高效的最近邻算法

Question

我有 10,00,000 个代理，每个代理都与 (x,y) 相关联坐标。我试图找到彼此接近的代理 ( radius=1.5 )。我尝试使用 PyTorch 来实现这一点:

X = torch.DoubleTensor(1000000,2).uniform_(0,10000)
torch.cdist(X,X,p=2)

但是，这样 session 就会崩溃。我在 google colab 上运行这个。当我尝试使用 radius_neighbors_graph 构建图形时发生了同样的情况。 scikit-learn 包。如果有人建议一种内存有效的方法来实现相同的功能，那将会有很大帮助。

Answer 1

您不太可能在不仔细考虑的情况下完整地计算 1M*1M 矩阵。你可能想要一些类似 scipy.spatial.KDTree 的东西.构建树后，您可以将代理的坐标传递给 query 方法来获取它在一定半径内的邻居。要一次获取所有邻居，您可以计算类似 sparse_distance_matrix 的内容。树本身处于适当的阈值。
或者，您可以研究任意数量的高效聚类算法。