Numpy - 寻找矩阵乘法的最近邻居
Numpy - Finding Nearest Neighbors of a Matrix Multiplication
我有一个包含 1000 个 128 维特征的数据集，例如形状为(1000,128).
我想找到形状为 (128,1) 的 128 维特征的排序最近邻。
通过数据集 (1000,128) 和特征 (128,1) 之间的矩阵乘法计算的距离将给出 (1000,1) 形状的相似数组：
数据集 (1000,128) x 特征 (128,1) = 相似度 (1000,1)
这是通过以下方式完成的：
# features.shape=(1000,128) ; feature.shape=(128,1) ; similarities.shape=(1000,1)
similarities = features.dot(feature)
计算距离（相似度）后，我使用以下代码查找最近的邻居：
# The n Nearest Neighbors Indexes (But Not Sorted)
nearest_neighbours_indexes_unsorted = np.argpartition(similarities, kth=-n)[-n:]
# The n Nearest Neighbors (But Not Sorted)
nearest_neighbours_similarities_unsorted = similarities[nearest_neighbours_indexes_unsorted]
# The Indexes of n Nearest Neighbors Sorted
nearest_neighbours_indexes_sorted = np.flip(nearest_neighbours_indexes_unsorted[np.argsort(nearest_neighbours_similarities_unsorted)], axis=0)
此代码对数百万数据的处理速度非常快（我很想知道是否有人可以提供更快的提示）但我希望能够一次性找到多个特征的最近邻居：
数据集 (1000,128) x 特征 (128,n) = 相似度 (1000,n)
一种方法是在一个循环中为每个特征计算上述代码（这很慢），另一种方法是更改代码以适应多维索引，这就是我遇到的问题：我没有知道如何为形状为 (128,n) 而不是 (128,1) 的特征编写上述代码。
帮助函数获取最大、最小的 n 索引、沿轴的元素
这是一个辅助函数，用于 select 顶部 n-largest 索引沿通用 ndarray 的通用轴，使用 np.argpartition and np.take_along_axis -
def take_largest_indices_along_axis(ar, n, axis):
s = ar.ndim*[slice(None,None,None)]
s[axis] = slice(-n,None,None)
idx = np.argpartition(ar, kth=-n, axis=axis)[tuple(s)]
sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis)
return np.flip(np.take_along_axis(idx, sidx, axis=axis),axis=axis)
扩展它以获得 n 个最小索引 -
def take_smallest_indices_along_axis(ar, n, axis):
s = ar.ndim*[slice(None,None,None)]
s[axis] = slice(None,n,None)
idx = np.argpartition(ar, kth=n, axis=axis)[tuple(s)]
sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis)
return np.take_along_axis(idx, sidx, axis=axis)
并将这些扩展到 select 最大或最小的 n 元素本身，只需使用 np.take_along_axis 的简单用法，如下所列 -
def take_largest_along_axis(ar, n, axis):
idx = take_largest_indices_along_axis(ar, n, axis)
return np.take_along_axis(ar, idx, axis=axis)
def take_smallest_along_axis(ar, n, axis):
idx = take_smallest_indices_along_axis(ar, n, axis)
return np.take_along_axis(ar, idx, axis=axis)
样本运行s
# Sample setup
In [200]: np.random.seed(0)
...: ar = np.random.randint(0,99,(5,5))
In [201]: ar
Out[201]:
array([[44, 47, 64, 67, 67],
[ 9, 83, 21, 36, 87],
[70, 88, 88, 12, 58],
[65, 39, 87, 46, 88],
[81, 37, 25, 77, 72]])
取最大 n 索引，元素沿轴 -
In [202]: take_largest_indices_along_axis(ar, n=2, axis=0)
Out[202]:
array([[4, 2, 2, 4, 3],
[2, 1, 3, 0, 1]])
In [203]: take_largest_indices_along_axis(ar, n=2, axis=1)
Out[203]:
array([[4, 3],
[4, 1],
[2, 1],
[4, 2],
[0, 3]])
In [251]: take_largest_along_axis(ar, n=2, axis=0)
Out[251]:
array([[81, 88, 88, 77, 88],
[70, 83, 87, 67, 87]])
In [252]: take_largest_along_axis(ar, n=2, axis=1)
Out[252]:
array([[67, 67],
[87, 83],
[88, 88],
[88, 87],
[81, 77]])
采用最小的 n 索引，元素沿轴 -
In [232]: take_smallest_indices_along_axis(ar, n=2, axis=0)
Out[232]:
array([[1, 4, 1, 2, 2],
[0, 3, 4, 1, 0]])
In [233]: take_smallest_indices_along_axis(ar, n=2, axis=1)
Out[233]:
array([[0, 1],
[0, 2],
[3, 4],
[1, 3],
[2, 1]])
In [253]: take_smallest_along_axis(ar, n=2, axis=0)
Out[253]:
array([[ 9, 37, 21, 12, 58],
[44, 39, 25, 36, 67]])
In [254]: take_smallest_along_axis(ar, n=2, axis=1)
Out[254]:
array([[44, 47],
[ 9, 21],
[12, 58],
[39, 46],
[25, 37]])
在这里破案
对于我们的案例，我们假设输入是 similarities 并且形状是 (1000,128)，表示 1000 个数据点和 128 个特征，我们想要寻找最大的 n=10 特征对于这些数据点中的每一个，它将是 -
take_largest_indices_along_axis(similarities, n=10, axis=1) # indices
take_largest_along_axis(similarities, n=10, axis=1) # elements
最终的 indices/values 数组的形状为 (1000, n)。
具有给定数据集形状的样本 运行 -
In [257]: np.random.seed(0)
...: similarities = np.random.randint(0,99,(1000,128))
In [263]: take_largest_indices_along_axis(similarities, n=10, axis=1).shape
Out[263]: (1000, 10)
In [264]: take_largest_along_axis(similarities, n=10, axis=1).shape
Out[264]: (1000, 10)
如果您希望为每个特征获取 n 最大数据点，即最终 indices/values 数组的形状为 (n, 128)，则使用 axis=0.
我有一个包含 1000 个 128 维特征的数据集，例如形状为(1000,128).
我想找到形状为 (128,1) 的 128 维特征的排序最近邻。
通过数据集 (1000,128) 和特征 (128,1) 之间的矩阵乘法计算的距离将给出 (1000,1) 形状的相似数组：
数据集 (1000,128) x 特征 (128,1) = 相似度 (1000,1)
这是通过以下方式完成的：
# features.shape=(1000,128) ; feature.shape=(128,1) ; similarities.shape=(1000,1)
similarities = features.dot(feature)
计算距离（相似度）后，我使用以下代码查找最近的邻居：
# The n Nearest Neighbors Indexes (But Not Sorted)
nearest_neighbours_indexes_unsorted = np.argpartition(similarities, kth=-n)[-n:]
# The n Nearest Neighbors (But Not Sorted)
nearest_neighbours_similarities_unsorted = similarities[nearest_neighbours_indexes_unsorted]
# The Indexes of n Nearest Neighbors Sorted
nearest_neighbours_indexes_sorted = np.flip(nearest_neighbours_indexes_unsorted[np.argsort(nearest_neighbours_similarities_unsorted)], axis=0)
此代码对数百万数据的处理速度非常快（我很想知道是否有人可以提供更快的提示）但我希望能够一次性找到多个特征的最近邻居：
数据集 (1000,128) x 特征 (128,n) = 相似度 (1000,n)
一种方法是在一个循环中为每个特征计算上述代码（这很慢），另一种方法是更改代码以适应多维索引，这就是我遇到的问题：我没有知道如何为形状为 (128,n) 而不是 (128,1) 的特征编写上述代码。
帮助函数获取最大、最小的 n 索引、沿轴的元素
这是一个辅助函数，用于 select 顶部 n-largest 索引沿通用 ndarray 的通用轴，使用 np.argpartition and np.take_along_axis -
def take_largest_indices_along_axis(ar, n, axis):
s = ar.ndim*[slice(None,None,None)]
s[axis] = slice(-n,None,None)
idx = np.argpartition(ar, kth=-n, axis=axis)[tuple(s)]
sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis)
return np.flip(np.take_along_axis(idx, sidx, axis=axis),axis=axis)
扩展它以获得 n 个最小索引 -
def take_smallest_indices_along_axis(ar, n, axis):
s = ar.ndim*[slice(None,None,None)]
s[axis] = slice(None,n,None)
idx = np.argpartition(ar, kth=n, axis=axis)[tuple(s)]
sidx = np.take_along_axis(ar,idx, axis=axis).argsort(axis=axis)
return np.take_along_axis(idx, sidx, axis=axis)
并将这些扩展到 select 最大或最小的 n 元素本身，只需使用 np.take_along_axis 的简单用法，如下所列 -
def take_largest_along_axis(ar, n, axis):
idx = take_largest_indices_along_axis(ar, n, axis)
return np.take_along_axis(ar, idx, axis=axis)
def take_smallest_along_axis(ar, n, axis):
idx = take_smallest_indices_along_axis(ar, n, axis)
return np.take_along_axis(ar, idx, axis=axis)
样本运行s
# Sample setup
In [200]: np.random.seed(0)
...: ar = np.random.randint(0,99,(5,5))
In [201]: ar
Out[201]:
array([[44, 47, 64, 67, 67],
[ 9, 83, 21, 36, 87],
[70, 88, 88, 12, 58],
[65, 39, 87, 46, 88],
[81, 37, 25, 77, 72]])
取最大 n 索引，元素沿轴 -
In [202]: take_largest_indices_along_axis(ar, n=2, axis=0)
Out[202]:
array([[4, 2, 2, 4, 3],
[2, 1, 3, 0, 1]])
In [203]: take_largest_indices_along_axis(ar, n=2, axis=1)
Out[203]:
array([[4, 3],
[4, 1],
[2, 1],
[4, 2],
[0, 3]])
In [251]: take_largest_along_axis(ar, n=2, axis=0)
Out[251]:
array([[81, 88, 88, 77, 88],
[70, 83, 87, 67, 87]])
In [252]: take_largest_along_axis(ar, n=2, axis=1)
Out[252]:
array([[67, 67],
[87, 83],
[88, 88],
[88, 87],
[81, 77]])
采用最小的 n 索引，元素沿轴 -
In [232]: take_smallest_indices_along_axis(ar, n=2, axis=0)
Out[232]:
array([[1, 4, 1, 2, 2],
[0, 3, 4, 1, 0]])
In [233]: take_smallest_indices_along_axis(ar, n=2, axis=1)
Out[233]:
array([[0, 1],
[0, 2],
[3, 4],
[1, 3],
[2, 1]])
In [253]: take_smallest_along_axis(ar, n=2, axis=0)
Out[253]:
array([[ 9, 37, 21, 12, 58],
[44, 39, 25, 36, 67]])
In [254]: take_smallest_along_axis(ar, n=2, axis=1)
Out[254]:
array([[44, 47],
[ 9, 21],
[12, 58],
[39, 46],
[25, 37]])
在这里破案
对于我们的案例，我们假设输入是 similarities 并且形状是 (1000,128)，表示 1000 个数据点和 128 个特征，我们想要寻找最大的 n=10 特征对于这些数据点中的每一个，它将是 -
take_largest_indices_along_axis(similarities, n=10, axis=1) # indices
take_largest_along_axis(similarities, n=10, axis=1) # elements
最终的 indices/values 数组的形状为 (1000, n)。
具有给定数据集形状的样本 运行 -
In [257]: np.random.seed(0)
...: similarities = np.random.randint(0,99,(1000,128))
In [263]: take_largest_indices_along_axis(similarities, n=10, axis=1).shape
Out[263]: (1000, 10)
In [264]: take_largest_along_axis(similarities, n=10, axis=1).shape
Out[264]: (1000, 10)
如果您希望为每个特征获取 n 最大数据点，即最终 indices/values 数组的形状为 (n, 128)，则使用 axis=0.