使用 pairwise_distances_chunked 计算最近邻搜索答案

【问题标题】：Using pairwise_distances_chunked to compute nearest neighbor search使用 pairwise_distances_chunked 计算最近邻搜索
【发布时间】：2021-11-15 03:12:56
【问题描述】：

我有一个细长的数据矩阵（大小：250,000 x 10），我将其表示为X。我还有一个向量p 来测量我的数据点的质量。我的目标是为我的数据矩阵X 中的每一行x 计算以下函数：

r(x) = min{ ||x-y|| | p[y]>p[x], y 在 X }

在较小的数据集上，我会使用 sklearn.metrics.pairwise_distances 来预先计算距离，如下所示：

from sklearn import metrics
n = len(X);

D_full = metrics.pairwise_distances(X);
r = np.zeros((n,1));
for i in range(n):
    r[i] = (D_full[i,p>p[i]]).min();

但是，上述方法会占用大量内存，因为我需要存储 D_full：一个完整的 n x n 矩阵。似乎sklearn.metrics.pairwise_distances_chunked 可能是解决此类问题的好工具，因为距离矩阵一次只存储一个块。我希望在如何使用它方面获得一些帮助，因为我目前不熟悉生成器对象。假设我调用以下代码：

from sklearn import metrics
D = metrics.pairwise_distances_chunked(X);
D_chunk = next(D)

上面的代码产生D（一个生成器对象）和D_chunk（一个536 x n 数组）。 D_chunk是否对应于我之前方法中矩阵D_full 的前536行？如果是，next(D_chunk) 是否对应接下来的 536 行？

非常感谢您的帮助。

【问题讨论】：

标签： python scikit-learn distance nearest-neighbor

【解决方案1】：

这是可能解决方案的概要，但缺少详细信息。简而言之，我会做以下事情：

创建一个BallTree 进行查询，并初始化大小为 250000 的 min_quality_distance，例如使用零。

为k=2

对于每个向量，找到最近的k 邻居（包括它自己）。
如果找到在k 内距离最远的向量，具有足够的质量，请针对该点更新min_quality_distance。
对于剩余的，重复k=k+1

在每次迭代中，我们必须查询更少的向量。这个想法是，在每次迭代中，您都会蚕食一些条件正确的最近邻居，并且每一步都会更容易。（简单 50%？）我将展示如何进行第一次迭代，这样应该可以构建循环。

你可以的；

import numpy as np
size = 250000

X = np.random.random( size=(size,10))
p = np.random.random( size=size)

然后用

创建一个 BallTree

from sklearn.neighbors import BallTree

tree = BallTree(X, leaf_size=10, metric='minkowski')

并查询它的第一次迭代（这大约需要 5 分钟。）

k_nearest = 2

distances, indici = tree.query(X, k=k_nearest, return_distance=True, dualtree=False, sort_results=True)

最近k内最远点的标记是

most_far_away_indici = indici[:,-1:]

它的质量

p[most_far_away_indici]

所以我们可以

quality_closeby = p[most_far_away_indici]

并检查它是否足够

indici_sufficient_quality = quality_closeby > np.expand_dims(p, axis=1)

我们有

found_closeby = np.all( indici_sufficient_quality, axis=1 )

这是真的，我们在附近找到了足够的质量。

我们可以更新向量

distances_nearby = distances[:,-1:]

rx = np.zeros(size)
rx[found_closeby] = distances_nearby[found_closeby][:,0]

我们现在需要照顾剩下的不走运的地方，这些是

~found_closeby

所以

indici_not_found = ~found_closeby

和

distances, indici = tree.query(X[indici_not_found], k=3, return_distance=True, dualtree=False, sort_results=True)

等等。

我确信前几个循环需要几分钟，但经过几次迭代后，速度会很快达到几秒钟。

这是一个与np.argwhere() 等的小练习，以确保正确的索引得到更新。

这可能不是最快的，但它是一种可行的方法。

【讨论】：

【解决方案2】：

由于无法知道某个块的尺寸，我建议使用 np.ones_like 而不是 np.zeros。

【讨论】：