np.argpartition 难以解决topK
topK是常用的一个功能,在python中,numpy等计算库使用了丰富的底层优化,对于矩阵计算的效率远高于python的for-loop实现。因此,我们希望尽量用一些numpy函数的组合实现topK。
pytorch
库提供了topk
函数,可以将高维数组沿某一维度(该维度共N项),选出最大(最小)的K项并排序。返回排序结果和index信息。奇怪的是,更轻量级的numpy库并没有直接提供 topK
函数。numpy只提供了argpartition
和 partition
,可以将最大(最小)的K项排到前K位。以argpartition
为例,最小的3项排到了前3位:
>>> x = np.array([3, 5, 6, 4, 2, 7, 1])
>>> x[np.argpartition(x, 3)]
array([2, 1, 3, 4, 5, 7, 6])
注意,argpartition
实现的是 partial sorting
,如上例,前3项和其余项被分开,但是两部分各自都是不排序的!而我们可能更想要topK的几项排好序(其余项则不作要求)。因此,下面提供一种基于argpartition的topK方法。
一个naive方法
最简单的方法自然是全排序,然后取前K项。缺点在于,要把topK之外的数据也进行排序,当K << N时较为浪费时间,复杂度为 O ( n log n ) O(n \log n) O(nlogn):
def naive_arg_topK(matrix, K, axis=0):
"""
perform topK based on np.argsort
:param matrix: to be sorted
:param K: select and sort the top K items
:param axis: dimension to be sorted.
:return:
"""
full_sort = np.argsort(matrix, axis=axis)
return full_sort.take(np.arange(K), axis=axis)
# Example
>>> dists = np.random.permutation(np.arange(30)).reshape(6, 5)
array([[17, 28, 1, 24, 23, 8],
[ 9, 21, 3, 22, 4, 5],
[19, 12, 26, 11, 13, 27],
[10, 15, 18, 14, 7, 16],
[ 0, 25, 29, 2, 6, 20]])
>>> naive_arg_topK(dists, 2, axis=0)
array([[4, 2, 0, 4, 1, 1],
[1, 3, 1, 2, 4, 0]])
>>> naive_arg_topK(dists, 2, axis=1)
array([[2, 5],
[2, 4],
[3, 1],
[4, 0],
[0, 3]])
基于partition的方法
对于 np.argpartition
函数,复杂度可能下降到
O
(
n
log
K
)
O(n \log K)
O(nlogK),很多情况下,K << N,此时naive方法有优化的空间。以下方法首先选出 topK 项,然后仅对前topK项进行排序(matrix仅限2d-array)。
def partition_arg_topK(matrix, K, axis=0):
"""
perform topK based on np.argpartition
:param matrix: to be sorted
:param K: select and sort the top K items
:param axis: 0 or 1. dimension to be sorted.
:return:
"""
a_part = np.argpartition(matrix, K, axis=axis)
if axis == 0:
row_index = np.arange(matrix.shape[1 - axis])
a_sec_argsort_K = np.argsort(matrix[a_part[0:K, :], row_index], axis=axis)
return a_part[0:K, :][a_sec_argsort_K, row_index]
else:
column_index = np.arange(matrix.shape[1 - axis])[:, None]
a_sec_argsort_K = np.argsort(matrix[column_index, a_part[:, 0:K]], axis=axis)
return a_part[:, 0:K][column_index, a_sec_argsort_K]
# Example
>>> dists = np.random.permutation(np.arange(30)).reshape(6, 5)
array([[17, 28, 1, 24, 23, 8],
[ 9, 21, 3, 22, 4, 5],
[19, 12, 26, 11, 13, 27],
[10, 15, 18, 14, 7, 16],
[ 0, 25, 29, 2, 6, 20]])
>>> partition_arg_topK(dists, 2, axis=0)
array([[4, 2, 0, 4, 1, 1],
[1, 3, 1, 2, 4, 0]])
>>> partition_arg_topK(dists, 2, axis=1)
array([[2, 5],
[2, 4],
[3, 1],
[4, 0],
[0, 3]])
大数据量测试
对shape(5000, 100000)的矩阵进行topK排序,测试时间为:
K | partition(s) | naive(s) |
---|---|---|
10 | 8.884 | 22.604 |
100 | 9.012 | 22.458 |
1000 | 8.904 | 22.506 |
5000 | 11.305 | 22.844 |
基本符合 log N / log K \log N / \log K logN/logK的比值。