numpy argsort 性能低下

Question

如果这个问题看起来很基础，我深表歉意。

给定：

>>> import numpy as np
>>> import time
>>> A = np.random.rand( int(1e5), int(5e4) ) # large numpy array

目标：

>>> bt=time.time(); B=np.argsort(A,axis=1);et=time.time();print(f"Took {(et-bt):.2f} s")

然而，计算索引数组需要相当长的时间：

# Took 316.90 s

问题：

还有其他节省时间的方法吗？

干杯，

Answer 1

输入数组 A 的形状为 (100_000, 50_000)，默认包含 np.float64 个值。这意味着您仅需要 8 * 100_000 * 50_000 / 1024**3 = 37.2 Gio 内存用于此数组。对于输出矩阵 B（应该包含 np.int64 类型的项），您可能还需要相同数量的 space。这意味着您需要一台至少有 74.4 Gio 的机器，更不用说操作系统 (OS) 和运行软件所需的 space（因此可能至少有 80 Gio）。如果你没有这样的内存space，那么你的OS将使用你的存储设备作为交换内存，这会慢得多。

假设你有这样的内存space可用，这样的计算是非常昂贵的。这主要是由于 页面错误 在填充 B 数组时，还有 B 相当大以及 Numpy 的默认实现计算顺序。您可以使用 并行 Numba 代码 和 较小的输出 来加快计算速度。这是一个例子：

import numba as nb

@nb.njit('int32[:,::1](float64[:,::1])', parallel=True)
def fastSort(a):
    b = np.empty(a.shape, dtype=np.int32)
    for i in nb.prange(a.shape[0]):
        b[i,:] = np.argsort(a[i,:])
    return b

请注意，Numba 的 argsort 实现效率低于 Numpy，但如果目标计算机具有 mny 个内核和良好的内存带宽，并行版本应该更快。

这是我的 6 核机器在大小为 (10_000, 50_000) 的矩阵上的结果（小 10 倍，因为我没有 80 Gio 的 RAM）：

Original implementation: 28.14 s
Sequential Numba:        38.70 s
Parallel Numba:           6.79 s

由此产生的解决方案因此快 4.1 倍。

请注意，在这种特定情况下，您甚至可以为 B 的项目使用类型 uint16，因为每行的大小小于 2**16 = 65536。这可能不会明显更快，但它应该节省一些额外的内存。由此产生的所需内存将为 46.5 Gio。您可以使用 np.float32 类型进一步减少所需的内存量（通常以损失 accuracy 为代价）。

如果您想进一步缩短执行时间，则需要使用 C 或 C++ 等低级语言针对您的特定需求更快地实现 argsort。但请注意，如果您不是使用这种语言的经验丰富的程序员或不熟悉低级优化，那么击败 Numpy 绝非易事。如果您对这样的解决方案感兴趣，一个好的开始可能是实施 基数排序.

numpy argsort 性能低下

numpy argsort slow performance

python

performance

numpy

numpy-ndarray